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

'use client';

import { useState, useRef, useEffect } from 'react';
import { Header } from '@/components/header';
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 } from '@/components/ui/card';
import { apiClient, type JobInfo, type ModelInfo } from '@/lib/api';
import { Loader2, Download, X, Upload } from 'lucide-react';
import { downloadImage, downloadAuthenticatedImage, fileToBase64 } from '@/lib/utils';
import { useLocalStorage } from '@/lib/hooks';

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

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

export default function UpscalerPage() {
  const [formData, setFormData] = useLocalStorage<UpscalerFormData>(
    'upscaler-form-data',
    defaultFormData
  );

  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 [previewImage, setPreviewImage] = useState<string | null>(null);
  const fileInputRef = useRef<HTMLInputElement>(null);
  const [upscalerModels, setUpscalerModels] = useState<ModelInfo[]>([]);

  useEffect(() => {
    const loadModels = async () => {
      try {
        // Fetch ESRGAN and upscaler models
        const [esrganModels, upscalerMods] = await Promise.all([
          apiClient.getModels('esrgan'),
          apiClient.getModels('upscaler'),
        ]);
        const allModels = [...esrganModels.models, ...upscalerMods.models];
        setUpscalerModels(allModels);
        // Set first model as default
        if (allModels.length > 0 && !formData.model) {
          setFormData(prev => ({ ...prev, model: allModels[0].name }));
        }
      } catch (err) {
        console.error('Failed to load upscaler models:', err);
      }
    };
    loadModels();
  }, []);

  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);
      setFormData((prev) => ({ ...prev, image: base64 }));
      setPreviewImage(base64);
      setError(null);
    } catch (err) {
      setError('Failed to load image');
    }
  };

  const pollJobStatus = async (jobId: string) => {
    const maxAttempts = 300;
    let attempts = 0;

    const poll = async () => {
      try {
        const status = await apiClient.getJobStatus(jobId);
        setJobInfo(status);

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

          // Handle both old format (result.images) and new format (outputs)
          if (status.outputs && status.outputs.length > 0) {
            // New format: convert output URLs to authenticated image URLs with cache-busting
            imageUrls = status.outputs.map((output: any) => {
              const filename = output.filename;
              return apiClient.getImageUrl(jobId, filename);
            });
          } else if (status.result?.images && status.result.images.length > 0) {
            // Old format: convert image URLs to authenticated URLs
            imageUrls = status.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);
        } else if (status.status === 'failed') {
          setError(status.error || 'Upscaling failed');
          setLoading(false);
        } else if (status.status === 'cancelled') {
          setError('Upscaling was cancelled');
          setLoading(false);
        } else if (attempts < maxAttempts) {
          attempts++;
          setTimeout(poll, 2000);
        } else {
          setError('Job polling timeout');
          setLoading(false);
        }
      } catch (err) {
        setError(err instanceof Error ? err.message : 'Failed to check job status');
        setLoading(false);
      }
    };

    poll();
  };

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

    if (!formData.image) {
      setError('Please upload an image first');
      return;
    }

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

    try {
      // 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`,
        // Add upscale-specific parameters here based on your API
      } as any);
      setJobInfo(job);
      const jobId = job.request_id || job.id;
      if (jobId) {
        await pollJobStatus(jobId);
      } 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');
      } 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="Source"
                          className="w-full rounded-lg border border-border"
                        />
                      </div>
                    )}
                    <Button
                      type="button"
                      variant="outline"
                      onClick={() => fileInputRef.current?.click()}
                      className="w-full"
                    >
                      <Upload className="h-4 w-4" />
                      {previewImage ? 'Change Image' : 'Upload Image'}
                    </Button>
                    <input
                      ref={fileInputRef}
                      type="file"
                      accept="image/*"
                      onChange={handleImageUpload}
                      className="hidden"
                    />
                  </div>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="upscale_factor">Upscale Factor</Label>
                  <select
                    id="upscale_factor"
                    name="upscale_factor"
                    value={formData.upscale_factor}
                    onChange={handleInputChange}
                    className="flex h-10 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>
                  <p className="text-xs text-muted-foreground">
                    Higher factors take longer to process
                  </p>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="model">Upscaling Model</Label>
                  <select
                    id="model"
                    name="model"
                    value={formData.model}
                    onChange={handleInputChange}
                    className="flex h-10 w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                  >
                    {upscalerModels.length > 0 ? (
                      upscalerModels.map((model) => (
                        <option key={model.id} value={model.name}>
                          {model.name}
                        </option>
                      ))
                    ) : (
                      <option value="">Loading models...</option>
                    )}
                  </select>
                  {upscalerModels.length === 0 && !loading && (
                    <p className="text-xs text-yellow-600 dark:text-yellow-400">
                      No upscaler models found. Please add ESRGAN or upscaler models to your models directory.
                    </p>
                  )}
                </div>

                <div className="flex gap-2">
                  <Button type="submit" disabled={loading || !formData.image} 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>
                )}

                {loading && jobInfo && (
                  <div className="rounded-md bg-muted p-3 text-sm">
                    <p>Job ID: {jobInfo.id || jobInfo.request_id || 'N/A'}</p>
                    <p>Status: {jobInfo.status}</p>
                    {jobInfo.progress !== undefined && (
                      <p>Progress: {Math.round(jobInfo.progress * 100)}%</p>
                    )}
                  </div>
                )}

                <div className="rounded-md bg-blue-500/10 p-3 text-sm text-blue-600 dark:text-blue-400">
                  <p className="font-medium">Note</p>
                  <p className="mt-1">
                    Upscaling functionality depends on your backend configuration and available upscaler models.
                  </p>
                </div>
              </form>
            </CardContent>
          </Card>

          {/* Right Panel - Upscaled Images */}
          <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>
  );
}