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

'use client';

import { useState, 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 { Textarea } from '@/components/ui/textarea';
import { PromptTextarea } from '@/components/prompt-textarea';
import { Label } from '@/components/ui/label';
import { Card, CardContent } from '@/components/ui/card';
import { apiClient, type GenerationRequest, type JobInfo, type ModelInfo } from '@/lib/api';
import { Loader2, Download, X, Trash2, RotateCcw, Power } from 'lucide-react';
import { downloadImage, downloadAuthenticatedImage } from '@/lib/utils';
import { useLocalStorage } from '@/lib/hooks';
import { ModelSelectionProvider, useModelSelection, useCheckpointSelection, useModelTypeSelection } from '@/contexts/model-selection-context';
import { EnhancedModelSelect, EnhancedModelSelectGroup } from '@/components/enhanced-model-select';
import { ModelSelectionWarning, AutoSelectionStatus } from '@/components/model-selection-indicator';

const defaultFormData: GenerationRequest = {
  prompt: '',
  negative_prompt: '',
  width: 512,
  height: 512,
  steps: 20,
  cfg_scale: 7.5,
  seed: '',
  sampling_method: 'euler_a',
  scheduler: 'default',
  batch_count: 1,
};

function Text2ImgForm() {
  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<GenerationRequest>(
    'text2img-form-data',
    defaultFormData
  );

  const [loading, setLoading] = useState(false);
  const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
  const [generatedImages, setGeneratedImages] = useState<string[]>([]);
  const [samplers, setSamplers] = useState<Array<{ name: string; description: string }>>([]);
  const [schedulers, setSchedulers] = useState<Array<{ name: string; description: string }>>([]);
  const [loraModels, setLoraModels] = useState<string[]>([]);
  const [embeddings, setEmbeddings] = useState<string[]>([]);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    const loadOptions = async () => {
      try {
        const [samplersData, schedulersData, modelsData, loras, embeds] = await Promise.all([
          apiClient.getSamplers(),
          apiClient.getSchedulers(),
          apiClient.getModels(), // Get all models with enhanced info
          apiClient.getModels('lora'),
          apiClient.getModels('embedding'),
        ]);
        setSamplers(samplersData);
        setSchedulers(schedulersData);
        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 options:', err);
      }
    };
    loadOptions();
  }, [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' || name === 'scheduler'
        ? value
        : Number(value),
    }));
  };

  const pollJobStatus = async (jobId: string) => {
    const maxAttempts = 300; // 5 minutes with 1 second interval
    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 || 'Generation failed');
          setLoading(false);
        } else if (status.status === 'cancelled') {
          setError('Generation 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 handleGenerate = async (e: React.FormEvent) => {
    e.preventDefault();

    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,
        model: selectedCheckpoint || undefined,
        vae: selectedVae || undefined,
      };

      const job = await apiClient.text2img(requestData);
      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 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');
      } catch (err) {
        console.error('Failed to cancel job:', err);
      }
    }
  };

  const handleClearPrompts = () => {
    setFormData({ ...formData, prompt: '', negative_prompt: '' });
  };

  const handleResetToDefaults = () => {
    setFormData(defaultFormData);
    setSelectedCheckpoint(null);
    setSelectedVae('');
    actions.resetSelection();
  };

  const handleServerRestart = async () => {
    if (!confirm('Are you sure you want to restart the server? This will cancel all running jobs.')) {
      return;
    }
    try {
      setLoading(true);
      await apiClient.restartServer();
      setError('Server restart initiated. Please wait...');
      setTimeout(() => {
        window.location.reload();
      }, 3000);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to restart server');
      setLoading(false);
    }
  };

  return (
    <AppLayout>
      <Header title="Text to Image" description="Generate images from 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 htmlFor="prompt">Prompt *</Label>
                  <PromptTextarea
                    value={formData.prompt}
                    onChange={(value) => setFormData({ ...formData, prompt: value })}
                    placeholder="a beautiful landscape with mountains and a lake, sunset, highly detailed..."
                    rows={4}
                    loras={loraModels}
                    embeddings={embeddings}
                  />
                  <p className="text-xs text-muted-foreground">
                    Tip: Use &lt;lora:name:weight&gt; for LoRAs (e.g., &lt;lora:myLora:0.8&gt;) 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="blurry, low quality, distorted..."
                    rows={2}
                    loras={loraModels}
                    embeddings={embeddings}
                  />
                </div>

                {/* Utility Buttons */}
                <div className="flex gap-2">
                  <Button
                    type="button"
                    variant="outline"
                    size="sm"
                    onClick={handleClearPrompts}
                    disabled={loading}
                    title="Clear both prompts"
                  >
                    <Trash2 className="h-4 w-4 mr-1" />
                    Clear Prompts
                  </Button>
                  <Button
                    type="button"
                    variant="outline"
                    size="sm"
                    onClick={handleResetToDefaults}
                    disabled={loading}
                    title="Reset all fields to defaults"
                  >
                    <RotateCcw className="h-4 w-4 mr-1" />
                    Reset to Defaults
                  </Button>
                  <Button
                    type="button"
                    variant="outline"
                    size="sm"
                    onClick={handleServerRestart}
                    disabled={loading}
                    title="Restart the backend server"
                  >
                    <Power className="h-4 w-4 mr-1" />
                    Restart Server
                  </Button>
                </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"
                  >
                    {samplers.length > 0 ? (
                      samplers.map((sampler) => (
                        <option key={sampler.name} value={sampler.name}>
                          {sampler.name.toUpperCase()} - {sampler.description}
                        </option>
                      ))
                    ) : (
                      <option value="euler_a">Loading...</option>
                    )}
                  </select>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="scheduler">Scheduler</Label>
                  <select
                    id="scheduler"
                    name="scheduler"
                    value={formData.scheduler}
                    onChange={handleInputChange}
                    className="flex h-10 w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                  >
                    {schedulers.length > 0 ? (
                      schedulers.map((scheduler) => (
                        <option key={scheduler.name} value={scheduler.name}>
                          {scheduler.name.toUpperCase()} - {scheduler.description}
                        </option>
                      ))
                    ) : (
                      <option value="default">Loading...</option>
                    )}
                  </select>
                </div>

                {/* Model Selection Section */}
                <EnhancedModelSelectGroup
                  title="Model Selection"
                  description="Select the checkpoint and additional models for generation"
                >
                  {/* 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 */}
                  <EnhancedModelSelect
                    modelType="vae"
                    label="VAE Model"
                    description="Optional VAE model for improved image quality"
                    value={selectedVae}
                    availableModels={vaeModels}
                    isAutoSelected={isVaeAutoSelected}
                    isUserOverride={isVaeUserOverride}
                    isLoading={isAutoSelecting}
                    onValueChange={setSelectedVae}
                    onSetUserOverride={setVaeUserOverride}
                    onClearOverride={clearVaeUserOverride}
                    placeholder="Use default VAE"
                  />

                  {/* Auto-selection Status */}
                  <div className="pt-2">
                    <AutoSelectionStatus
                      isAutoSelecting={isAutoSelecting}
                      hasAutoSelection={Object.keys(state.autoSelectedModels).length > 0}
                    />
                  </div>

                  {/* Warnings and Errors */}
                  <ModelSelectionWarning
                    warnings={warnings}
                    errors={error ? [error] : []}
                    onClearWarnings={actions.clearWarnings}
                  />
                </EnhancedModelSelectGroup>

                <div className="space-y-2">
                  <Label htmlFor="batch_count">Batch Count</Label>
                  <Input
                    id="batch_count"
                    name="batch_count"
                    type="number"
                    value={formData.batch_count}
                    onChange={handleInputChange}
                    min={1}
                    max={4}
                  />
                </div>

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

                {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>
                )}
              </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, `generated-${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, `generated-${Date.now()}-${index}.png`);
                              });
                          }}
                        >
                          <Download className="h-4 w-4" />
                        </Button>
                      </div>
                    ))}
                  </div>
                )}
              </div>
            </CardContent>
          </Card>
        </div>
      </div>
    </AppLayout>
  );
}

export default function Text2ImgPage() {
  return <Text2ImgForm />;
}