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

"use client";

import { useState, useEffect, useRef } 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 { Textarea } from "@/components/ui/textarea";
import { PromptTextarea } from "@/components/forms";
import { Label } from "@/components/ui/label";
import {
  Card,
  CardContent,
  CardHeader,
  CardTitle,
  CardDescription,
} from "@/components/ui/card";
import { InpaintingCanvas } from "@/components/features/image-generation";
import { apiClient, type JobInfo, type JobDetailsResponse } from "@/lib/api";
import { Loader2, X, Download } from "lucide-react";
import { downloadAuthenticatedImage } from "@/lib/utils";
import { useLocalStorage } from "@/lib/storage";
import {
  ModelSelectionProvider,
  useModelSelection,
  useCheckpointSelection,
  useModelTypeSelection,
} from "@/contexts/model-selection-context";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
// import { AutoSelectionStatus } from '@/components/features/models';

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

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

function InpaintingForm() {
  const { state, actions } = useModelSelection();
  const {
    checkpointModels,
    selectedCheckpointModel,
    selectedCheckpoint,
    setSelectedCheckpoint,
    isAutoSelecting,
    warnings,
    error: checkpointError,
  } = useCheckpointSelection();

  const {
    availableModels: vaeModels,
    selectedModel: selectedVae,
    isUserOverride: isVaeUserOverride,
    isAutoSelected: isVaeAutoSelected,
    setSelectedModel: setSelectedVae,
    setUserOverride: setVaeUserOverride,
    clearUserOverride: clearVaeUserOverride,
  } = useModelTypeSelection("vae");

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

  // Separate state for image data (not stored in localStorage)
  const [sourceImage, setSourceImage] = useState<string>("");
  const [maskImage, setMaskImage] = useState<string>("");

  const [loading, setLoading] = useState(false);
  const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
  const [generatedImages, setGeneratedImages] = useState<string[]>([]);
  const [loraModels, setLoraModels] = useState<string[]>([]);
  const [embeddings, setEmbeddings] = useState<string[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);

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

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

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

  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 handleSourceImageChange = (image: string) => {
    setSourceImage(image);
    setError(null);
  };

  const handleMaskImageChange = (image: string) => {
    setMaskImage(image);
    setError(null);
  };

  const pollJobStatus = async (jobId: string) => {
    const maxAttempts = 300; // 5 minutes with 2 second interval
    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);
                }
              }
              return imageUrl; // Already a full URL
            });
          }

          // 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 (!sourceImage) {
      setError("Please upload a source image first");
      return;
    }

    if (!maskImage) {
      setError("Please create a mask first");
      return;
    }

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

    try {
      // Validate model selection
      if (selectedCheckpointModel) {
        const validation = actions.validateSelection(selectedCheckpointModel);
        if (!validation.isValid) {
          setError(
            `Missing required models: ${validation.missingRequired.join(", ")}`,
          );
          setLoading(false);
          return;
        }
      }

      const requestData = {
        ...formData,
        source_image: sourceImage,
        mask_image: maskImage,
        model: selectedCheckpoint || undefined,
        vae: selectedVae || undefined,
      };

      const job = await apiClient.inpainting(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="Inpainting"
        description="Edit images by masking areas and regenerating with AI"
      />
      <div className="container mx-auto p-6">
        <div className="grid gap-6 lg:grid-cols-2">
          {/* Left Panel - Canvas and Form */}
          <div className="space-y-6">
            <InpaintingCanvas
              onSourceImageChange={handleSourceImageChange}
              onMaskImageChange={handleMaskImageChange}
              targetWidth={formData.width}
              targetHeight={formData.height}
            />

            <Card>
              <CardContent className="pt-6">
                <form onSubmit={handleGenerate} className="space-y-4">
                  <div className="space-y-2">
                    <Label htmlFor="prompt">Prompt *</Label>
                    <PromptTextarea
                      value={formData.prompt}
                      onChange={(value) =>
                        setFormData({ ...formData, prompt: value })
                      }
                      placeholder="Describe what to generate in the masked areas..."
                      rows={3}
                      loras={loraModels}
                      embeddings={embeddings}
                    />
                    <p className="text-xs text-muted-foreground">
                      Tip: Use {"<lora:name:weight>"} 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 in the generated areas..."
                      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}
                      />
                    </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}
                      />
                    </div>
                  </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>

                  {/* Model Selection Section */}
                  <Card>
                    <CardHeader>
                      <CardTitle>Model Selection</CardTitle>
                      <CardDescription>
                        Select checkpoint and additional models for generation
                      </CardDescription>
                    </CardHeader>
                    {/* Checkpoint Selection */}
                    <div className="space-y-2">
                      <Label htmlFor="checkpoint">Checkpoint Model *</Label>
                      <select
                        id="checkpoint"
                        value={selectedCheckpoint || ""}
                        onChange={(e) =>
                          setSelectedCheckpoint(e.target.value || null)
                        }
                        className="flex h-10 w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                        disabled={isAutoSelecting}
                      >
                        <option value="">Select a checkpoint model...</option>
                        {checkpointModels.map((model) => (
                          <option key={model.id} value={model.name}>
                            {model.name} {model.loaded ? "(Loaded)" : ""}
                          </option>
                        ))}
                      </select>
                    </div>

                    {/* VAE Selection */}
                    <div className="space-y-2">
                      <Label>VAE Model (Optional)</Label>
                      <Select
                        value={selectedVae || ""}
                        onValueChange={(value) => {
                          if (value) {
                            setSelectedVae(value);
                            setVaeUserOverride(value);
                          } else {
                            clearVaeUserOverride();
                          }
                        }}
                        disabled={isAutoSelecting}
                      >
                        <SelectTrigger>
                          <SelectValue placeholder="Use default VAE" />
                        </SelectTrigger>
                        <SelectContent>
                          <SelectItem value="">Use default VAE</SelectItem>
                          {vaeModels.map((model) => (
                            <SelectItem key={model.name} value={model.name}>
                              {model.name}
                            </SelectItem>
                          ))}
                        </SelectContent>
                      </Select>
                      {isVaeAutoSelected && (
                        <p className="text-xs text-muted-foreground">
                          Auto-selected VAE model
                        </p>
                      )}
                    </div>
                  </Card>

                  <div className="flex gap-2">
                    <Button
                      type="submit"
                      disabled={loading || !sourceImage || !maskImage}
                      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>
          </div>

          {/* 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,
                              `inpainting-${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>
    </AppLayout>
  );
}

export default function InpaintingPage() {
  return <InpaintingForm />;
}