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

"use client";

import { useState, useRef, useEffect } from "react";
import { Header, AppLayout } from "@/components/layout";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Textarea } from "@/components/ui/textarea";
import { PromptTextarea } from "@/components/forms";
import { Label } from "@/components/ui/label";
import { Card, CardContent } from "@/components/ui/card";
import { ImageInput } from "@/components/ui/image-input";
import { apiClient, type JobInfo, type JobDetailsResponse } from "@/lib/api";
import {
  downloadImage,
  downloadAuthenticatedImage,
  fileToBase64,
} from "@/lib/utils";
import { useLocalStorage, useMemoryStorage } from "@/lib/storage";
import { type ImageValidationResult } from "@/lib/image-validation";
import { Loader2, Download, X } from "lucide-react";

type Img2ImgFormData = {
  prompt: string;
  negative_prompt: string;
  image: string;
  strength: number;
  steps: number;
  cfg_scale: number;
  seed: string;
  sampling_method: string;
  width?: number;
  height?: number;
};

const defaultFormData: Img2ImgFormData = {
  prompt: "",
  negative_prompt: "",
  image: "",
  strength: 0.75,
  steps: 20,
  cfg_scale: 7.5,
  seed: "",
  sampling_method: "euler_a",
  width: 512,
  height: 512,
};

function Img2ImgForm() {
  // Store form data without the image to avoid localStorage quota issues
  const { image: _, ...formDataWithoutImage } = defaultFormData;
  const [formData, setFormData] = useLocalStorage<
    Omit<Img2ImgFormData, "image">
  >("img2img-form-data", formDataWithoutImage);

  // Store image separately in memory
  const [imageData, setImageData] = useMemoryStorage<string>("");

  // Combined form data with image
  const fullFormData = { ...formData, image: imageData };

  const [loading, setLoading] = useState(false);
  const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
  const [generatedImages, setGeneratedImages] = useState<string[]>([]);
  const [previewImage, setPreviewImage] = useState<string | null>(null);
  const [loraModels, setLoraModels] = useState<string[]>([]);
  const [embeddings, setEmbeddings] = useState<string[]>([]);
  const [selectedImage, setSelectedImage] = useState<File | string | null>(
    null,
  );
  const [imageValidation, setImageValidation] =
    useState<ImageValidationResult | null>(null);
  const [originalImage, setOriginalImage] = useState<string | null>(null);
  const [isResizing, setIsResizing] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
  const [imageLoadingFromUrl, setImageLoadingFromUrl] = useState(false);

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

  useEffect(() => {
    const loadModels = async () => {
      try {
        const [loras, embeds] = await Promise.all([
          apiClient.getModels("lora"),
          apiClient.getModels("embedding"),
        ]);
        setLoraModels(loras.models.map((m) => m.name));
        setEmbeddings(embeds.models.map((m) => m.name));
      } catch (err) {
        console.error("Failed to load models:", err);
      }
    };
    loadModels();
  }, []);

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

  const handleImageChange = async (image: File | string | null) => {
    setSelectedImage(image);
    setError(null);

    if (!image) {
      setImageData("");
      setFormData((prev) => ({ ...prev, image: "" }));
      setPreviewImage(null);
      setImageValidation(null);
      setOriginalImage(null);
      return;
    }

    try {
      let imageBase64: string;
      let previewUrl: string | null;

      if (image instanceof File) {
        // Convert File to base64
        imageBase64 = await fileToBase64(image);
        previewUrl = imageBase64;
      } else {
        // For URLs, don't set preview immediately - wait for validation
        // The validation will provide base64 data for preview and processing
        imageBase64 = ""; // Will be set by validation
        previewUrl = null; // Will be set by validation
        setImageLoadingFromUrl(true);
        console.log(
          "Image URL provided, waiting for validation to process:",
          image,
        );
      }

      // Store original image for resizing (will be updated by validation if URL)
      if (image instanceof File) {
        setOriginalImage(imageBase64);
        setImageData(imageBase64);
        setPreviewImage(previewUrl);
      }
      setFormData((prev) => ({ ...prev })); // Don't store image in localStorage
    } catch (err) {
      setError("Failed to process image");
      console.error("Image processing error:", err);
    }
  };

  // Auto-resize image when width or height changes, but only if we have valid image data
  useEffect(() => {
    const resizeImage = async () => {
      if (!originalImage || !formData.width || !formData.height) {
        return;
      }

      // Don't resize if we're already resizing or still loading from URL
      if (isResizing || imageLoadingFromUrl) {
        return;
      }

      // Check if we have valid image data (data URL or HTTP URL)
      const isValidImageData =
        originalImage.startsWith("data:image/") ||
        originalImage.startsWith("http://") ||
        originalImage.startsWith("https://");

      if (!isValidImageData) {
        console.warn(
          "Invalid image data format for resizing:",
          originalImage.substring(0, 50),
        );
        return; // Don't show error for timing issues, just skip resize
      }

      try {
        setIsResizing(true);

        console.log(
          "Attempting to resize image, data type:",
          originalImage.startsWith("data:image/") ? "data URL" : "HTTP URL",
        );
        console.log("Image data length:", originalImage.length);
        console.log("Resize dimensions:", formData.width, "x", formData.height);

        const result = await apiClient.resizeImage(
          originalImage,
          formData.width,
          formData.height,
        );
        setImageData(result.image);
        setFormData((prev) => ({ ...prev })); // Don't store image in localStorage
        setPreviewImage(result.image);
      } catch (err) {
        console.error("Failed to resize image:", err);
        const errorMessage =
          err instanceof Error ? err.message : "Unknown error";
        setError(
          `Failed to resize image: ${errorMessage}. You may need to resize the image manually or use different dimensions.`,
        );
      } finally {
        setIsResizing(false);
      }
    };

    resizeImage();
  }, [
    formData.width,
    formData.height,
    originalImage,
    imageLoadingFromUrl,
    isResizing,
  ]);

  const handleImageValidation = (result: ImageValidationResult) => {
    setImageValidation(result);
    setImageLoadingFromUrl(false);
    if (!result.isValid) {
      setError(result.error || "Invalid image");
    } else {
      setError(null);
      // If we have temporary URL or base64 data from URL download, use it for preview and processing
      if (selectedImage && typeof selectedImage === "string") {
        if (result.tempUrl) {
          // Use temporary URL for preview and processing
          const fullTempUrl = result.tempUrl.startsWith("/")
            ? `${window.location.origin}${result.tempUrl}`
            : result.tempUrl;

          setPreviewImage(fullTempUrl);
          setOriginalImage(fullTempUrl);

          // For processing, we still need to convert to base64 or use the URL directly
          // The resize endpoint can handle URLs, so we can use the temp URL
          setImageData(fullTempUrl);

          console.log(
            "Using temporary URL from URL validation for image processing:",
            fullTempUrl,
          );
          console.log("Temporary filename:", result.tempFilename);
        } else if (result.base64Data) {
          // Fallback to base64 data if no temporary URL
          if (
            result.base64Data.startsWith("data:image/") &&
            result.base64Data.includes("base64,")
          ) {
            setPreviewImage(result.base64Data);
            setImageData(result.base64Data);
            setOriginalImage(result.base64Data);
            console.log(
              "Using base64 data from URL validation for image processing",
            );
            console.log("Base64 data length:", result.base64Data.length);
            console.log(
              "Base64 data preview:",
              result.base64Data.substring(0, 100) + "...",
            );
          } else {
            console.error(
              "Invalid base64 data format received from server:",
              result.base64Data.substring(0, 100),
            );
            setError("Invalid image data format received from server");
          }
        }
      }
    }
  };

  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 || "Generation failed");
          setLoading(false);
          isPolling = false;
        } else if (status.job.status === "cancelled") {
          setError("Generation 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 handleGenerate = async (e: React.FormEvent) => {
    e.preventDefault();

    if (!fullFormData.image) {
      setError("Please upload or select an image first");
      return;
    }

    // Check if image validation passed
    if (imageValidation && !imageValidation.isValid) {
      setError("Please fix the image validation errors before generating");
      return;
    }

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

    try {
      const requestData = {
        ...fullFormData,
      };

      const job = await apiClient.img2img(requestData);
      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 generate image");
      setLoading(false);
    }
  };

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

  return (
    <AppLayout>
      <Header
        title="Image to Image"
        description="Transform images with AI using text prompts"
      />
      <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={handleGenerate} className="space-y-4">
                <div className="space-y-2">
                  <Label>Source Image *</Label>
                  <ImageInput
                    value={selectedImage}
                    onChange={handleImageChange}
                    onValidation={handleImageValidation}
                    disabled={loading}
                    maxSize={10 * 1024 * 1024} // 10MB
                    accept="image/*"
                    placeholder="Enter image URL or select a file"
                    showPreview={true}
                  />
                </div>

                <div className="space-y-2">
                  <Label htmlFor="prompt">Prompt *</Label>
                  <PromptTextarea
                    value={formData.prompt}
                    onChange={(value) =>
                      setFormData({ ...formData, prompt: value })
                    }
                    placeholder="Describe the transformation you want..."
                    rows={3}
                    loras={loraModels}
                    embeddings={embeddings}
                  />
                  <p className="text-xs text-muted-foreground">
                    Tip: Use &lt;lora:name:weight&gt; for LoRAs and embedding
                    names directly
                  </p>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="negative_prompt">Negative Prompt</Label>
                  <PromptTextarea
                    value={formData.negative_prompt || ""}
                    onChange={(value) =>
                      setFormData({ ...formData, negative_prompt: value })
                    }
                    placeholder="What to avoid..."
                    rows={2}
                    loras={loraModels}
                    embeddings={embeddings}
                  />
                </div>

                <div className="space-y-2">
                  <Label htmlFor="strength">
                    Strength: {formData.strength.toFixed(2)}
                  </Label>
                  <Input
                    id="strength"
                    name="strength"
                    type="range"
                    value={formData.strength}
                    onChange={handleInputChange}
                    min={0}
                    max={1}
                    step={0.05}
                  />
                  <p className="text-xs text-muted-foreground">
                    Lower values preserve more of the original image
                  </p>
                </div>

                <div className="grid grid-cols-2 gap-4">
                  <div className="space-y-2">
                    <Label htmlFor="width">Width</Label>
                    <Input
                      id="width"
                      name="width"
                      type="number"
                      value={formData.width}
                      onChange={handleInputChange}
                      step={64}
                      min={256}
                      max={2048}
                      disabled={isResizing}
                    />
                  </div>
                  <div className="space-y-2">
                    <Label htmlFor="height">Height</Label>
                    <Input
                      id="height"
                      name="height"
                      type="number"
                      value={formData.height}
                      onChange={handleInputChange}
                      step={64}
                      min={256}
                      max={2048}
                      disabled={isResizing}
                    />
                  </div>
                </div>
                {isResizing && (
                  <div className="text-sm text-muted-foreground flex items-center gap-2">
                    <Loader2 className="h-4 w-4 animate-spin" />
                    Resizing image...
                  </div>
                )}

                <div className="grid grid-cols-2 gap-4">
                  <div className="space-y-2">
                    <Label htmlFor="steps">Steps</Label>
                    <Input
                      id="steps"
                      name="steps"
                      type="number"
                      value={formData.steps}
                      onChange={handleInputChange}
                      min={1}
                      max={150}
                    />
                  </div>
                  <div className="space-y-2">
                    <Label htmlFor="cfg_scale">CFG Scale</Label>
                    <Input
                      id="cfg_scale"
                      name="cfg_scale"
                      type="number"
                      value={formData.cfg_scale}
                      onChange={handleInputChange}
                      step={0.5}
                      min={1}
                      max={30}
                    />
                  </div>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="seed">Seed (optional)</Label>
                  <Input
                    id="seed"
                    name="seed"
                    value={formData.seed}
                    onChange={handleInputChange}
                    placeholder="Leave empty for random"
                  />
                </div>

                <div className="space-y-2">
                  <Label htmlFor="sampling_method">Sampling Method</Label>
                  <select
                    id="sampling_method"
                    name="sampling_method"
                    value={formData.sampling_method}
                    onChange={handleInputChange}
                    className="flex h-10 w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                  >
                    <option value="euler">Euler</option>
                    <option value="euler_a">Euler A</option>
                    <option value="heun">Heun</option>
                    <option value="dpm2">DPM2</option>
                    <option value="dpm++2s_a">DPM++ 2S A</option>
                    <option value="dpm++2m">DPM++ 2M</option>
                    <option value="dpm++2mv2">DPM++ 2M V2</option>
                    <option value="lcm">LCM</option>
                  </select>
                </div>

                <div className="flex gap-2">
                  <Button
                    type="submit"
                    disabled={
                      loading ||
                      !imageData ||
                      imageValidation?.isValid === false
                    }
                    className="flex-1"
                  >
                    {loading ? (
                      <>
                        <Loader2 className="h-4 w-4 animate-spin" />
                        Generating...
                      </>
                    ) : (
                      "Generate"
                    )}
                  </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>

          {/* Right Panel - Generated Images */}
          <Card>
            <CardContent className="pt-6">
              <div className="space-y-4">
                <h3 className="text-lg font-semibold">Generated Images</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
                        ? "Generating..."
                        : "Generated 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={`Generated ${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,
                              `img2img-${Date.now()}-${index}.png`,
                              authToken || undefined,
                              unixUser || undefined,
                            ).catch((err) => {
                              console.error("Failed to download image:", err);
                              // Fallback to regular download if authenticated download fails
                              downloadImage(
                                image,
                                `img2img-${Date.now()}-${index}.png`,
                              );
                            });
                          }}
                        >
                          <Download className="h-4 w-4" />
                        </Button>
                      </div>
                    ))}
                  </div>
                )}
              </div>
            </CardContent>
          </Card>
        </div>
      </div>
    </AppLayout>
  );
}

export default function Img2ImgPage() {
  return <Img2ImgForm />;
}