stable-diffusion.cpp-rest/webui/app/upscaler/page.tsx

"use client";

import { useState, useRef, useEffect, Suspense } from "react";
import { useSearchParams } from "next/navigation";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
  Card,
  CardContent,
} from "@/components/ui/card";
import {
  apiClient,
  type JobInfo,
  type JobDetailsResponse,
  type ModelInfo,
  type EnhancedModelsResponse,
} from "@/lib/api";
import { Loader2, Download, X, Upload } from "lucide-react";
import {
  downloadAuthenticatedImage,
  fileToBase64,
} from "@/lib/utils";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { AppLayout, Header } from "@/components/layout";

type UpscalerFormData = {
  upscale_factor: number;
  model: string;
};

const defaultFormData: UpscalerFormData = {
  upscale_factor: 2,
  model: "",
};



function UpscalerForm() {
  const searchParams = useSearchParams();

  // Simple state management - no complex hooks initially
  const [formData, setFormData] = useState<UpscalerFormData>(defaultFormData);

  // Separate state for image data (not stored in localStorage)
  const [uploadedImage, setUploadedImage] = useState<string>("");
  const [previewImage, setPreviewImage] = useState<string | null>(null);

  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
  const [generatedImages, setGeneratedImages] = useState<string[]>([]);
  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
  const fileInputRef = useRef<HTMLInputElement>(null);

  // URL input state
  const [urlInput, setUrlInput] = useState('');

  // Local state for upscaler models - no global context to avoid performance issues
  const [upscalerModels, setUpscalerModels] = useState<ModelInfo[]>([]);
  const [modelsLoading, setModelsLoading] = useState(false);

  // Cleanup polling on unmount
  useEffect(() => {
    return () => {
      if (pollCleanup) {
        pollCleanup();
      }
    };
  }, [pollCleanup]);

  // Load image from URL parameter on mount
  useEffect(() => {
    const imageUrl = searchParams.get('imageUrl');
    if (imageUrl) {
      loadImageFromUrl(imageUrl);
    }
  }, [searchParams]);

  // Load upscaler models on mount
  useEffect(() => {
    let isComponentMounted = true;

    const loadModels = async () => {
      try {
        setModelsLoading(true);
        setError(null);

        // Set up timeout for API call
        const timeoutPromise = new Promise((_, reject) =>
          setTimeout(() => reject(new Error('API call timeout')), 5000)
        );

        const apiPromise = apiClient.getModels("esrgan");
        const modelsData = await Promise.race([apiPromise, timeoutPromise]) as EnhancedModelsResponse;

        console.log("API call completed, models:", modelsData.models?.length || 0);

        if (!isComponentMounted) return;

        // Set models locally - no global state updates
        setUpscalerModels(modelsData.models || []);

        // Set first model as default if none selected
        if (modelsData.models?.length > 0 && !formData.model) {
          setFormData((prev) => ({
            ...prev,
            model: modelsData.models[0].name,
          }));
        }
      } catch (err) {
        console.error("Failed to load upscaler models:", err);
        if (isComponentMounted) {
          setError(`Failed to load upscaler models: ${err instanceof Error ? err.message : 'Unknown error'}`);
        }
      } finally {
        if (isComponentMounted) {
          setModelsLoading(false);
        }
      }
    };

    loadModels();

    return () => {
      isComponentMounted = false;
    };
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  const loadImageFromUrl = async (url: string) => {
    try {
      setError(null);
      // Fetch the image and convert to base64
      const response = await fetch(url);
      if (!response.ok) {
        throw new Error('Failed to fetch image');
      }
      const blob = await response.blob();
      const base64 = await new Promise<string>((resolve, reject) => {
        const reader = new FileReader();
        reader.onload = () => resolve(reader.result as string);
        reader.onerror = reject;
        reader.readAsDataURL(blob);
      });
      setUploadedImage(base64);
      setPreviewImage(base64);
    } catch (err) {
      console.error('Failed to load image from URL:', err);
      setError('Failed to load image from gallery');
    }
  };

  const handleImageUpload = async (e: React.ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0];
    if (!file) return;

    try {
      const base64 = await fileToBase64(file);
      setUploadedImage(base64);
      setPreviewImage(base64);
      setError(null);
    } catch {
      setError("Failed to load image");
    }
  };

  const pollJobStatus = async (jobId: string) => {
    const maxAttempts = 300;
    let attempts = 0;
    let isPolling = true;
    let timeoutId: NodeJS.Timeout | null = null;

    const poll = async () => {
      if (!isPolling) return;

      try {
        const status: JobDetailsResponse = await apiClient.getJobStatus(jobId);
        setJobInfo(status.job);

        if (status.job.status === "completed") {
          let imageUrls: string[] = [];

          // Handle both old format (result.images) and new format (outputs)
          if (status.job.outputs && status.job.outputs.length > 0) {
            // New format: convert output URLs to authenticated image URLs with cache-busting
            imageUrls = status.job.outputs.map((output: { filename: string }) => {
              const filename = output.filename;
              return apiClient.getImageUrl(jobId, filename);
            });
          } else if (
            status.job.result?.images &&
            status.job.result.images.length > 0
          ) {
            // Old format: convert image URLs to authenticated URLs
            imageUrls = status.job.result.images.map((imageUrl: string) => {
              // Extract filename from URL if it's already a full URL
              if (imageUrl.includes("/output/")) {
                const parts = imageUrl.split("/output/");
                if (parts.length === 2) {
                  const filename = parts[1].split("?")[0]; // Remove query params
                  return apiClient.getImageUrl(jobId, filename);
                }
              }
              // If it's just a filename, convert it directly
              return apiClient.getImageUrl(jobId, imageUrl);
            });
          }

          // Create a new array to trigger React re-render
          setGeneratedImages([...imageUrls]);
          setLoading(false);
          isPolling = false;
        } else if (status.job.status === "failed") {
          setError(status.job.error_message || status.job.error || "Upscaling failed");
          setLoading(false);
          isPolling = false;
        } else if (status.job.status === "cancelled") {
          setError("Upscaling was cancelled");
          setLoading(false);
          isPolling = false;
        } else if (attempts < maxAttempts) {
          attempts++;
          timeoutId = setTimeout(poll, 2000);
        } else {
          setError("Job polling timeout");
          setLoading(false);
          isPolling = false;
        }
      } catch {
        if (isPolling) {
          setError("Failed to check job status");
          setLoading(false);
          isPolling = false;
        }
      }
    };

    poll();

    // Return cleanup function
    return () => {
      isPolling = false;
      if (timeoutId) {
        clearTimeout(timeoutId);
      }
    };
  };

  const handleUpscale = async (e: React.FormEvent) => {
    e.preventDefault();

    if (!uploadedImage) {
      setError("Please upload an image first");
      return;
    }

    setLoading(true);
    setError(null);
    setGeneratedImages([]);
    setJobInfo(null);

    try {
      // Validate model selection
      if (!formData.model) {
        setError("Please select an upscaler model");
        setLoading(false);
        return;
      }

      // Unload all currently loaded models and load the selected upscaler model
      const selectedModel = upscalerModels.find(m => m.name === formData.model);
      const modelId = selectedModel?.id || selectedModel?.sha256;
      if (!selectedModel) {
        setError("Selected upscaler model not found.");
        setLoading(false);
        return;
      }
      if (!modelId) {
        setError("Selected upscaler model does not have a hash. Please compute the hash on the models page.");
        setLoading(false);
        return;
      }

      try {
        // Get all loaded models
        const loadedModels = await apiClient.getAllModels(undefined, true);

        // Unload all loaded models
        for (const model of loadedModels) {
          const unloadId = model.id || model.sha256;
          if (unloadId) {
            try {
              await apiClient.unloadModel(unloadId);
            } catch (unloadErr) {
              console.warn(`Failed to unload model ${model.name}:`, unloadErr);
              // Continue with others
            }
          }
        }

        // Load the selected upscaler model
        await apiClient.loadModel(modelId);
      } catch (modelErr) {
        console.error("Failed to prepare upscaler model:", modelErr);
        setError("Failed to prepare upscaler model. Please try again.");
        setLoading(false);
        return;
      }

      const job = await apiClient.upscale({
        image: uploadedImage,
        model: formData.model,
        upscale_factor: formData.upscale_factor,
      });
      setJobInfo(job);
      const jobId = job.request_id || job.id;
      if (jobId) {
        const cleanup = pollJobStatus(jobId);
        setPollCleanup(() => cleanup);
      } else {
        setError("No job ID returned from server");
        setLoading(false);
      }
    } catch {
      setError("Failed to upscale image");
      setLoading(false);
    }
  };

  const handleCancel = async () => {
    const jobId = jobInfo?.request_id || jobInfo?.id;
    if (jobId) {
      try {
        await apiClient.cancelJob(jobId);
        setLoading(false);
        setError("Upscaling cancelled");
        // Cleanup polling
        if (pollCleanup) {
          pollCleanup();
          setPollCleanup(null);
        }
      } catch (err) {
        console.error("Failed to cancel job:", err);
      }
    }
  };

  return (
    <AppLayout>
      <Header
        title="Upscaler"
        description="Enhance and upscale your images with AI"
      />
      <div className="container mx-auto p-6">
        <div className="grid gap-6 lg:grid-cols-2">
          {/* Left Panel - Form Parameters */}
          <div className="space-y-6">
            <Card>
              <CardContent className="pt-6">
                <form onSubmit={handleUpscale} className="space-y-4">
                  {/* Image Upload Section */}
                  <div className="space-y-2">
                    <Label htmlFor="image-upload">Image *</Label>
                    <div className="space-y-2">
                      <input
                        id="image-upload"
                        type="file"
                        accept="image/*"
                        onChange={handleImageUpload}
                        ref={fileInputRef}
                        className="hidden"
                      />
                      <Button
                        type="button"
                        variant="outline"
                        onClick={() => fileInputRef.current?.click()}
                        className="w-full"
                      >
                        <Upload className="mr-2 h-4 w-4" />
                        Choose Image File
                      </Button>

                      <div className="flex gap-2">
                        <Input
                          type="url"
                          placeholder="Or paste image URL"
                          value={urlInput}
                          onChange={(e) => setUrlInput(e.target.value)}
                          className="flex-1"
                        />
                        <Button
                          type="button"
                          variant="outline"
                          onClick={() => loadImageFromUrl(urlInput)}
                          disabled={!urlInput}
                        >
                          Load
                        </Button>
                      </div>
                    </div>
                  </div>

                  {/* Model Selection */}
                  <div className="space-y-2">
                    <Label htmlFor="model">Upscaler Model</Label>
                    <Select
                      value={formData.model}
                      onValueChange={(value) =>
                        setFormData((prev) => ({ ...prev, model: value }))
                      }
                    >
                      <SelectTrigger>
                        <SelectValue placeholder="Select upscaler model" />
                      </SelectTrigger>
                      <SelectContent>
                        {upscalerModels.map((model) => (
                          <SelectItem key={model.name} value={model.name}>
                            {model.name}
                          </SelectItem>
                        ))}
                      </SelectContent>
                    </Select>
                    {modelsLoading && (
                      <p className="text-sm text-muted-foreground">Loading models...</p>
                    )}
                  </div>

                  {/* Upscale Factor */}
                  <div className="space-y-2">
                    <Label htmlFor="upscale_factor">
                      Upscale Factor *
                      <span className="text-xs text-muted-foreground ml-1">
                        ({formData.upscale_factor}x)
                      </span>
                    </Label>
                    <input
                      id="upscale_factor"
                      name="upscale_factor"
                      type="range"
                      min="2"
                      max="8"
                      step="1"
                      value={formData.upscale_factor}
                      onChange={(e) =>
                        setFormData((prev) => ({
                          ...prev,
                          upscale_factor: Number(e.target.value),
                        }))
                      }
                      className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer dark:bg-gray-700"
                    />
                    <div className="flex justify-between text-xs text-muted-foreground">
                      <span>2x</span>
                      <span>8x</span>
                    </div>
                  </div>

                  {/* Generate Button */}
                  <Button type="submit" disabled={loading || !uploadedImage} className="w-full">
                    {loading ? (
                      <>
                        <Loader2 className="mr-2 h-4 w-4 animate-spin" />
                        Upscaling...
                      </>
                    ) : (
                      "Upscale Image"
                    )}
                  </Button>

                  {/* Cancel Button */}
                  {jobInfo && (
                    <Button
                      type="button"
                      variant="outline"
                      onClick={handleCancel}
                      disabled={!loading}
                      className="w-full"
                    >
                      <X className="h-4 w-4 mr-2" />
                      Cancel
                    </Button>
                  )}

                  {/* Error Display */}
                  {error && (
                    <div className="rounded-md bg-destructive/10 p-3 text-sm text-destructive">
                      {error}
                    </div>
                  )}
                </form>
              </CardContent>
            </Card>
          </div>

          {/* Right Panel - Image Preview and Results */}
          <div className="space-y-6">
            {/* Image Preview */}
            <Card>
              <CardContent className="pt-6">
                <div className="space-y-4">
                  <h3 className="text-lg font-semibold">Image Preview</h3>
                  {previewImage ? (
                    <div className="relative">
                      <img
                        src={previewImage}
                        alt="Preview"
                        className="w-full rounded-lg border border-border"
                      />
                    </div>
                  ) : (
                    <div className="flex h-64 items-center justify-center rounded-lg border-2 border-dashed border-border">
                      <p className="text-muted-foreground">
                        Upload an image to see preview
                      </p>
                    </div>
                  )}
                </div>
              </CardContent>
            </Card>

            {/* Results */}
            <Card>
              <CardContent className="pt-6">
                <div className="space-y-4">
                  <h3 className="text-lg font-semibold">Upscaled Images</h3>
                  {generatedImages.length === 0 ? (
                    <div className="flex h-64 items-center justify-center rounded-lg border-2 border-dashed border-border">
                      <p className="text-muted-foreground">
                        {loading
                          ? "Upscaling in progress..."
                          : "Upscaled images will appear here"}
                      </p>
                    </div>
                  ) : (
                    <div className="grid gap-4">
                      {generatedImages.map((image, index) => (
                        <div key={index} className="relative group">
                          <img
                            src={image}
                            alt={`Upscaled ${index + 1}`}
                            className="w-full rounded-lg border border-border"
                          />
                          <Button
                            size="icon"
                            variant="secondary"
                            className="absolute top-2 right-2 opacity-0 group-hover:opacity-100 transition-opacity"
                            onClick={() => {
                              const authToken =
                                localStorage.getItem("auth_token");
                              const unixUser = localStorage.getItem("unix_user");
                              downloadAuthenticatedImage(
                                image,
                                `upscaled-${Date.now()}-${index}.png`,
                                authToken || undefined,
                                unixUser || undefined,
                              ).catch((err) => {
                                console.error("Failed to download image:", err);
                              });
                            }}
                          >
                            <Download className="h-4 w-4" />
                          </Button>
                        </div>
                      ))}
                    </div>
                  )}
                </div>
              </CardContent>
            </Card>
          </div>
        </div>
      </div>
    </AppLayout>
  );
}

export default function UpscalerPage() {
  return (
    <Suspense fallback={<div>Loading...</div>}>
      <UpscalerForm />
    </Suspense>
  );
}