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

"use client";

import { useState, useRef, useEffect } from "react";
import { Header } from "@/components/layout";
import { AppLayout } from "@/components/layout";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
  Card,
  CardContent,
  CardHeader,
  CardTitle,
  CardDescription,
} from "@/components/ui/card";
import {
  apiClient,
  type JobInfo,
  type ModelInfo,
  type JobDetailsResponse,
} from "@/lib/api";
import { Loader2, Download, X, Upload } from "lucide-react";
import {
  downloadImage,
  downloadAuthenticatedImage,
  fileToBase64,
} from "@/lib/utils";
import { useLocalStorage } from "@/lib/storage";
import {
  ModelSelectionProvider,
  useModelSelection,
  useModelTypeSelection,
} from "@/contexts/model-selection-context";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
// import { AutoSelectionStatus } from '@/components/features/models';

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

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

function UpscalerForm() {
  const { state, actions } = useModelSelection();

  const {
    availableModels: upscalerModels,
    selectedModel: selectedUpscalerModel,
    isUserOverride: isUpscalerUserOverride,
    isAutoSelected: isUpscalerAutoSelected,
    setSelectedModel: setSelectedUpscalerModel,
    setUserOverride: setUpscalerUserOverride,
    clearUserOverride: clearUpscalerUserOverride,
  } = useModelTypeSelection("upscaler");

  const [formData, setFormData] = useLocalStorage<UpscalerFormData>(
    "upscaler-form-data",
    defaultFormData,
    { excludeLargeData: true, maxSize: 512 * 1024 },
  );

  // 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);

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

  useEffect(() => {
    const loadModels = async () => {
      try {
        // Fetch all models with enhanced info
        const modelsData = await apiClient.getModels();
        // Filter for upscaler models (ESRGAN and upscaler types)
        const allUpscalerModels = [
          ...modelsData.models.filter((m) => m.type.toLowerCase() === "esrgan"),
          ...modelsData.models.filter(
            (m) => m.type.toLowerCase() === "upscaler",
          ),
        ];
        actions.setModels(modelsData.models);

        // Set first model as default if none selected
        if (allUpscalerModels.length > 0 && !formData.model) {
          setFormData((prev) => ({
            ...prev,
            model: allUpscalerModels[0].name,
          }));
        }
      } catch (err) {
        console.error("Failed to load upscaler models:", err);
      }
    };
    loadModels();
  }, [actions, formData.model, setFormData]);

  // Update form data when upscaler model changes
  useEffect(() => {
    if (selectedUpscalerModel) {
      setFormData((prev) => ({
        ...prev,
        model: selectedUpscalerModel,
      }));
    }
  }, [selectedUpscalerModel, setFormData]);

  const handleInputChange = (
    e: React.ChangeEvent<HTMLInputElement | HTMLSelectElement>,
  ) => {
    const { name, value } = e.target;
    setFormData((prev) => ({
      ...prev,
      [name]: name === "upscale_factor" ? Number(value) : value,
    }));
  };

  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 (err) {
      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: any) => {
              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 || "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 (err) {
        if (isPolling) {
          setError(
            err instanceof Error ? err.message : "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 (!selectedUpscalerModel) {
        setError("Please select an upscaler model");
        setLoading(false);
        return;
      }

      // Note: You may need to adjust the API endpoint based on your backend implementation
      const job = await apiClient.generateImage({
        prompt: `upscale ${formData.upscale_factor}x`,
        model: selectedUpscalerModel,
        image: uploadedImage,
        // Add upscale-specific parameters here based on your API
      } as any);
      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 (err) {
      setError(err instanceof Error ? err.message : "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 */}
          <Card>
            <CardContent className="pt-6">
              <form onSubmit={handleUpscale} className="space-y-4">
                <div className="space-y-2">
                  <Label>Source Image *</Label>
                  <div className="space-y-4">
                    {previewImage && (
                      <div className="relative">
                        <img
                          src={previewImage}
                          alt="Preview"
                          className="h-64 w-full rounded-lg object-cover"
                        />
                        <Button
                          type="button"
                          variant="destructive"
                          size="icon"
                          className="absolute top-2 right-2 h-8 w-8"
                          onClick={() => {
                            setPreviewImage(null);
                            setUploadedImage("");
                          }}
                        >
                          <X className="h-4 w-4" />
                        </Button>
                      </div>
                    )}

                    <div className="space-y-2">
                      <div className="flex items-center justify-center">
                        <input
                          ref={fileInputRef}
                          type="file"
                          accept="image/*"
                          onChange={handleImageUpload}
                          className="hidden"
                        />
                        <Button
                          type="button"
                          variant="outline"
                          onClick={() => fileInputRef.current?.click()}
                        >
                          <Upload className="mr-2 h-4 w-4" />
                          Choose Image
                        </Button>
                      </div>
                    </div>
                  </div>
                </div>

                <div className="space-y-2">
                  <Label>Upscaling Factor</Label>
                  <select
                    value={formData.upscale_factor}
                    onChange={(e) =>
                      setFormData((prev) => ({
                        ...prev,
                        upscale_factor: Number(e.target.value),
                      }))
                    }
                    className="w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                  >
                    <option value={2}>2x (Double)</option>
                    <option value={3}>3x (Triple)</option>
                    <option value={4}>4x (Quadruple)</option>
                  </select>
                </div>

                <div className="space-y-2">
                  <Label>Upscaler Model</Label>
                  <Select
                    value={formData.model}
                    onValueChange={(value) => {
                      setFormData((prev) => ({ ...prev, model: value }));
                      setSelectedUpscalerModel(value);
                      setUpscalerUserOverride(value);
                    }}
                  >
                    <SelectTrigger>
                      <SelectValue placeholder="Select an upscaler model" />
                    </SelectTrigger>
                    <SelectContent>
                      {upscalerModels.map((model) => (
                        <SelectItem key={model.name} value={model.name}>
                          {model.name}
                        </SelectItem>
                      ))}
                    </SelectContent>
                  </Select>
                </div>

                <div className="flex gap-2">
                  <Button
                    type="submit"
                    disabled={loading || !uploadedImage}
                    className="flex-1"
                  >
                    {loading ? (
                      <>
                        <Loader2 className="h-4 w-4 animate-spin" />
                        Upscaling...
                      </>
                    ) : (
                      "Upscale"
                    )}
                  </Button>
                  {loading && (
                    <Button
                      type="button"
                      variant="destructive"
                      onClick={handleCancel}
                    >
                      <X className="h-4 w-4" />
                      Cancel
                    </Button>
                  )}
                </div>

                {error && (
                  <div className="rounded-md bg-destructive/10 p-3 text-sm text-destructive">
                    {error}
                  </div>
                )}
              </form>
            </CardContent>
          </Card>
          <Card>
            <CardContent className="pt-6">
              <div className="space-y-4">
                <h3 className="text-lg font-semibold">Upscaled Image</h3>
                {generatedImages.length === 0 ? (
                  <div className="flex h-96 items-center justify-center rounded-lg border-2 border-dashed border-border">
                    <p className="text-muted-foreground">
                      {loading
                        ? "Upscaling..."
                        : "Upscaled image 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()}-${formData.upscale_factor}x.png`,
                              authToken || undefined,
                              unixUser || undefined,
                            ).catch((err) => {
                              console.error("Failed to download image:", err);
                              // Fallback to regular download if authenticated download fails
                              downloadImage(
                                image,
                                `upscaled-${Date.now()}-${formData.upscale_factor}x.png`,
                              );
                            });
                          }}
                        >
                          <Download className="h-4 w-4" />
                        </Button>
                      </div>
                    ))}
                  </div>
                )}
              </div>
            </CardContent>
          </Card>
        </div>
      </div>
    </AppLayout>
  );
}

export default function UpscalerPage() {
  return <UpscalerForm />;
}