469 lines
17 KiB
TypeScript
469 lines
17 KiB
TypeScript
'use client';
|
|
|
|
import { useState, useEffect } from 'react';
|
|
import { useRouter, useSearchParams } from 'next/navigation';
|
|
import { toast } from 'react-hot-toast';
|
|
import { ImagePlus, Download, Sparkles, Pencil, X, Loader2, Maximize, Film } from 'lucide-react';
|
|
import JobProgress from '@/components/JobProgress';
|
|
import ProviderControls from '@/components/ProviderControls';
|
|
import { modulesApi, assetsApi, capabilitiesApi } from '@/lib/api';
|
|
import { useStore } from '@/lib/store';
|
|
import { ProviderConfig } from '@/types/providers';
|
|
|
|
export default function ImageGeneratePage() {
|
|
const { addJob, updateJob } = useStore();
|
|
const router = useRouter();
|
|
const searchParams = useSearchParams();
|
|
|
|
// Provider config state
|
|
const [capabilities, setCapabilities] = useState<Record<string, ProviderConfig> | null>(null);
|
|
const [loadingCapabilities, setLoadingCapabilities] = useState(true);
|
|
|
|
// Generation state
|
|
const [prompt, setPrompt] = useState('');
|
|
const [provider, setProvider] = useState('openai');
|
|
const [model, setModel] = useState('');
|
|
const [providerOptions, setProviderOptions] = useState<Record<string, any>>({});
|
|
const [jobId, setJobId] = useState<string | null>(null);
|
|
const [generatedImages, setGeneratedImages] = useState<any[]>([]);
|
|
const [loading, setLoading] = useState(false);
|
|
|
|
// Iterative editing state (for Nano Banana)
|
|
const [editingImage, setEditingImage] = useState<any | null>(null);
|
|
const [editInstructions, setEditInstructions] = useState('');
|
|
|
|
// Load provider capabilities on mount
|
|
useEffect(() => {
|
|
const loadCapabilities = async () => {
|
|
try {
|
|
const response = await capabilitiesApi.getImageProviders();
|
|
setCapabilities(response.data);
|
|
|
|
// Set default provider and model
|
|
const firstProvider = Object.keys(response.data)[0];
|
|
setProvider(firstProvider);
|
|
setModel(response.data[firstProvider].defaultModel);
|
|
|
|
// Initialize with default values
|
|
initializeDefaults(response.data[firstProvider]);
|
|
} catch (err) {
|
|
console.error('Failed to load provider configurations:', err);
|
|
toast.error('Failed to load provider configurations');
|
|
} finally {
|
|
setLoadingCapabilities(false);
|
|
}
|
|
};
|
|
|
|
loadCapabilities();
|
|
}, []);
|
|
|
|
// Handle URL parameters
|
|
useEffect(() => {
|
|
const urlPrompt = searchParams.get('prompt');
|
|
if (urlPrompt) {
|
|
setPrompt(urlPrompt);
|
|
}
|
|
|
|
// Check for reference asset for editing/variations if we support it via URL
|
|
// (Optional: handle assetId if needed)
|
|
}, [searchParams]);
|
|
|
|
// Initialize default values for provider
|
|
const initializeDefaults = (config: ProviderConfig) => {
|
|
if (!config) {
|
|
console.error('Config is undefined');
|
|
return;
|
|
}
|
|
|
|
const defaults: Record<string, any> = {};
|
|
|
|
// Common controls
|
|
if (config.commonControls && Array.isArray(config.commonControls)) {
|
|
config.commonControls.forEach((control) => {
|
|
defaults[control.name] = control.default;
|
|
});
|
|
}
|
|
|
|
// Model-specific controls
|
|
const modelConfig = config.models?.find(m => m.id === config.defaultModel);
|
|
if (modelConfig?.controls && Array.isArray(modelConfig.controls)) {
|
|
modelConfig.controls.forEach((control) => {
|
|
defaults[control.name] = control.default;
|
|
});
|
|
}
|
|
|
|
setProviderOptions(defaults);
|
|
};
|
|
|
|
// Handle provider change
|
|
const handleProviderChange = (newProvider: string) => {
|
|
if (!capabilities) return;
|
|
|
|
setProvider(newProvider);
|
|
const config = capabilities[newProvider];
|
|
setModel(config.defaultModel);
|
|
initializeDefaults(config);
|
|
|
|
// Cancel editing if changing provider
|
|
if (editingImage) {
|
|
setEditingImage(null);
|
|
setEditInstructions('');
|
|
}
|
|
};
|
|
|
|
// Handle model change
|
|
const handleModelChange = (newModel: string) => {
|
|
setModel(newModel);
|
|
|
|
if (!capabilities) return;
|
|
const config = capabilities[provider];
|
|
const modelConfig = config.models.find(m => m.id === newModel);
|
|
|
|
// Merge current options with model defaults
|
|
const modelDefaults: Record<string, any> = {};
|
|
modelConfig?.controls?.forEach((control) => {
|
|
if (!(control.name in providerOptions)) {
|
|
modelDefaults[control.name] = control.default;
|
|
}
|
|
});
|
|
|
|
setProviderOptions({
|
|
...providerOptions,
|
|
...modelDefaults
|
|
});
|
|
};
|
|
|
|
const handleGenerate = async () => {
|
|
const effectivePrompt = editingImage ? editInstructions : prompt;
|
|
if (!effectivePrompt.trim()) {
|
|
toast.error(editingImage ? 'Please enter edit instructions' : 'Please enter a prompt');
|
|
return;
|
|
}
|
|
|
|
setLoading(true);
|
|
setGeneratedImages([]); // Always clear previous images
|
|
setJobId(null); // Reset job ID
|
|
|
|
try {
|
|
const payload = {
|
|
prompt: effectivePrompt,
|
|
provider: editingImage ? 'nano-banana' : provider,
|
|
model: editingImage ? 'gemini-2.5-flash-image' : model,
|
|
provider_options: editingImage ? undefined : providerOptions,
|
|
reference_asset_id: editingImage?.id || undefined,
|
|
};
|
|
|
|
// Debug logging
|
|
if (editingImage) {
|
|
console.log('🎨 Nano Banana Edit Mode:');
|
|
console.log(' Reference Asset ID:', editingImage.id);
|
|
console.log(' Edit Instructions:', effectivePrompt);
|
|
console.log(' Full Payload:', payload);
|
|
}
|
|
|
|
const response = await modulesApi.generateImage(payload);
|
|
|
|
const job = response.data;
|
|
setJobId(job.id);
|
|
addJob({
|
|
id: job.id,
|
|
module: 'image_generation',
|
|
status: job.status,
|
|
progress: job.progress,
|
|
created_at: job.created_at,
|
|
});
|
|
|
|
toast.success(editingImage ? 'Image editing started!' : 'Image generation started!');
|
|
} catch (err: any) {
|
|
console.error('Generation error:', err);
|
|
toast.error(err.response?.data?.detail || 'Failed to start generation');
|
|
setLoading(false);
|
|
}
|
|
};
|
|
|
|
const handleJobComplete = async (job: any) => {
|
|
setLoading(false);
|
|
updateJob(job.id, { status: 'completed', progress: 100 });
|
|
|
|
if (job.output_asset_ids?.length > 0) {
|
|
const images = await Promise.all(
|
|
job.output_asset_ids.map(async (id: string) => {
|
|
const asset = await assetsApi.get(id);
|
|
return asset.data;
|
|
})
|
|
);
|
|
|
|
// When editing, replace with new edited version to maintain chain visibility
|
|
if (editingImage) {
|
|
setGeneratedImages(images);
|
|
// Keep the editing image as the new base for next edit
|
|
setEditingImage(images[0]);
|
|
setEditInstructions('');
|
|
toast.success('Image edited successfully! You can continue editing.');
|
|
} else {
|
|
setGeneratedImages(images);
|
|
toast.success('Images generated successfully!');
|
|
}
|
|
}
|
|
};
|
|
|
|
const handleStartEdit = (image: any) => {
|
|
setEditingImage(image);
|
|
setEditInstructions('');
|
|
// Auto-switch to Nano Banana for editing
|
|
if (capabilities && capabilities['nano-banana']) {
|
|
setProvider('nano-banana');
|
|
setModel('gemini-2.5-flash-image');
|
|
initializeDefaults(capabilities['nano-banana']);
|
|
}
|
|
};
|
|
|
|
const handleCancelEdit = () => {
|
|
setEditingImage(null);
|
|
setEditInstructions('');
|
|
};
|
|
|
|
const handleJobError = (error: string) => {
|
|
setLoading(false);
|
|
toast.error(error);
|
|
};
|
|
|
|
const handleDownload = async (assetId: string, filename: string) => {
|
|
try {
|
|
const response = await assetsApi.download(assetId);
|
|
const url = window.URL.createObjectURL(response.data);
|
|
const a = document.createElement('a');
|
|
a.href = url;
|
|
a.download = filename;
|
|
a.click();
|
|
window.URL.revokeObjectURL(url);
|
|
} catch (err) {
|
|
toast.error('Failed to download image');
|
|
}
|
|
};
|
|
|
|
if (loadingCapabilities) {
|
|
return (
|
|
<div className="flex items-center justify-center min-h-screen">
|
|
<Loader2 className="w-8 h-8 animate-spin text-forge-yellow" />
|
|
</div>
|
|
);
|
|
}
|
|
|
|
const currentConfig = capabilities?.[provider];
|
|
const supportsEditing = provider === 'nano-banana' || provider === 'gemini';
|
|
|
|
return (
|
|
<div className="max-w-6xl mx-auto space-y-8">
|
|
<div className="flex items-center gap-4">
|
|
<div className="w-12 h-12 bg-forge-yellow/10 rounded-lg flex items-center justify-center">
|
|
<ImagePlus className="w-6 h-6 text-forge-yellow" />
|
|
</div>
|
|
<div>
|
|
<h1 className="text-2xl font-bold text-white">Image Generator</h1>
|
|
<p className="text-gray-500">Create stunning images with AI</p>
|
|
</div>
|
|
</div>
|
|
|
|
<div className="grid grid-cols-1 lg:grid-cols-2 gap-8">
|
|
{/* Controls */}
|
|
<div className="space-y-6">
|
|
{/* Editing Mode Panel */}
|
|
{editingImage && (
|
|
<div className="bg-purple-900/20 border border-purple-500/50 rounded-xl p-4 space-y-4">
|
|
<div className="flex items-center justify-between">
|
|
<h3 className="text-purple-400 font-medium flex items-center gap-2">
|
|
<Pencil className="w-4 h-4" />
|
|
Editing Image
|
|
</h3>
|
|
<button
|
|
onClick={handleCancelEdit}
|
|
className="text-gray-400 hover:text-white transition-colors"
|
|
>
|
|
<X className="w-5 h-5" />
|
|
</button>
|
|
</div>
|
|
<div className="flex gap-4">
|
|
<div className="w-24 h-24 rounded-lg overflow-hidden border border-purple-500/50">
|
|
<img
|
|
src={`/api/v1/assets/${editingImage.id}/download`}
|
|
alt="Reference"
|
|
className="w-full h-full object-cover"
|
|
/>
|
|
</div>
|
|
<div className="flex-1">
|
|
<label className="block text-sm font-medium text-gray-300 mb-2">
|
|
Edit Instructions
|
|
</label>
|
|
<textarea
|
|
value={editInstructions}
|
|
onChange={(e) => setEditInstructions(e.target.value)}
|
|
placeholder="Describe how you want to modify this image..."
|
|
className="input-field min-h-[80px] resize-none"
|
|
/>
|
|
</div>
|
|
</div>
|
|
<p className="text-xs text-gray-500">
|
|
Using Nano Banana (Gemini) for iterative editing
|
|
<br />
|
|
Reference: {editingImage.original_filename || editingImage.id.substring(0, 8)}
|
|
</p>
|
|
</div>
|
|
)}
|
|
|
|
{/* Prompt */}
|
|
{!editingImage && (
|
|
<div>
|
|
<label className="block text-sm font-medium text-gray-300 mb-2">
|
|
Prompt
|
|
</label>
|
|
<textarea
|
|
value={prompt}
|
|
onChange={(e) => setPrompt(e.target.value)}
|
|
placeholder="Describe the image you want to create..."
|
|
className="input-field min-h-[120px] resize-none"
|
|
/>
|
|
</div>
|
|
)}
|
|
|
|
{/* Provider & Model Selection */}
|
|
<div className="grid grid-cols-2 gap-4">
|
|
<div>
|
|
<label className="block text-sm font-medium text-gray-300 mb-2">
|
|
Provider
|
|
</label>
|
|
<select
|
|
value={provider}
|
|
onChange={(e) => handleProviderChange(e.target.value)}
|
|
className="select-field"
|
|
disabled={editingImage !== null}
|
|
>
|
|
{capabilities && Object.entries(capabilities).map(([id, config]) => (
|
|
<option key={id} value={id}>
|
|
{config.name}
|
|
</option>
|
|
))}
|
|
</select>
|
|
</div>
|
|
<div>
|
|
<label className="block text-sm font-medium text-gray-300 mb-2">
|
|
Model
|
|
</label>
|
|
<select
|
|
value={model}
|
|
onChange={(e) => handleModelChange(e.target.value)}
|
|
className="select-field"
|
|
disabled={editingImage !== null}
|
|
>
|
|
{currentConfig?.models.map((m) => (
|
|
<option key={m.id} value={m.id}>
|
|
{m.name}
|
|
</option>
|
|
))}
|
|
</select>
|
|
</div>
|
|
</div>
|
|
|
|
{/* Model Description */}
|
|
{currentConfig?.models.find(m => m.id === model)?.description && (
|
|
<p className="text-xs text-gray-500 -mt-4">
|
|
{currentConfig.models.find(m => m.id === model)?.description}
|
|
</p>
|
|
)}
|
|
|
|
{/* Dynamic Provider Controls */}
|
|
{currentConfig && !editingImage && (
|
|
<div className="bg-forge-gray rounded-lg p-4">
|
|
<ProviderControls
|
|
config={currentConfig}
|
|
selectedModel={model}
|
|
values={providerOptions}
|
|
onChange={setProviderOptions}
|
|
/>
|
|
</div>
|
|
)}
|
|
|
|
{/* Generate Button */}
|
|
<button
|
|
onClick={handleGenerate}
|
|
disabled={loading || (editingImage ? !editInstructions.trim() : !prompt.trim())}
|
|
className={`w-full flex items-center justify-center gap-2 disabled:opacity-50 disabled:cursor-not-allowed ${editingImage ? 'bg-purple-600 hover:bg-purple-700 text-white py-3 px-6 rounded-lg font-medium transition-colors' : 'btn-primary'
|
|
}`}
|
|
>
|
|
{editingImage ? <Pencil className="w-5 h-5" /> : <Sparkles className="w-5 h-5" />}
|
|
{loading ? (editingImage ? 'Editing...' : 'Generating...') : (editingImage ? 'Apply Edits' : 'Generate Images')}
|
|
</button>
|
|
|
|
{/* Job Progress */}
|
|
{jobId && loading && (
|
|
<JobProgress
|
|
jobId={jobId}
|
|
onComplete={handleJobComplete}
|
|
onError={handleJobError}
|
|
/>
|
|
)}
|
|
</div>
|
|
|
|
{/* Results */}
|
|
<div>
|
|
<h2 className="text-lg font-semibold text-white mb-4">Generated Images</h2>
|
|
{generatedImages.length > 0 ? (
|
|
<div className="space-y-4">
|
|
{generatedImages.map((image) => (
|
|
<div
|
|
key={image.id}
|
|
className="bg-forge-dark rounded-xl overflow-hidden border border-gray-800 group"
|
|
>
|
|
<div className="relative w-full group">
|
|
<img
|
|
src={`/api/v1/assets/${image.id}/download`}
|
|
alt="Generated"
|
|
className="w-full h-auto object-contain bg-black/20"
|
|
/>
|
|
<div className="absolute inset-0 bg-black/60 opacity-0 group-hover:opacity-100 transition-opacity flex flex-col items-center justify-center gap-2 p-4">
|
|
{supportsEditing && (
|
|
<button
|
|
onClick={() => handleStartEdit(image)}
|
|
className="bg-purple-600 hover:bg-purple-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
|
|
title="Edit with Nano Banana"
|
|
>
|
|
<Pencil className="w-4 h-4" />
|
|
Edit
|
|
</button>
|
|
)}
|
|
<button
|
|
onClick={() => router.push(`/image/upscale?assetId=${image.id}`)}
|
|
className="bg-blue-600 hover:bg-blue-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
|
|
>
|
|
<Maximize className="w-4 h-4" />
|
|
Upscale
|
|
</button>
|
|
<button
|
|
onClick={() => router.push(`/video/generate?assetId=${image.id}`)}
|
|
className="bg-green-600 hover:bg-green-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
|
|
>
|
|
<Film className="w-4 h-4" />
|
|
Video
|
|
</button>
|
|
<button
|
|
onClick={() => handleDownload(image.id, image.original_filename)}
|
|
className="btn-primary py-1.5 px-3 text-sm flex items-center gap-2 w-full justify-center"
|
|
>
|
|
<Download className="w-4 h-4" />
|
|
Download
|
|
</button>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
))}
|
|
</div>
|
|
) : (
|
|
<div className="bg-forge-dark rounded-xl border border-gray-800 aspect-square flex items-center justify-center">
|
|
<p className="text-gray-500">Generated images will appear here</p>
|
|
</div>
|
|
)}
|
|
</div>
|
|
</div>
|
|
</div>
|
|
);
|
|
}
|