forge/frontend/app/image/upscale/page.tsx

459 lines
17 KiB
TypeScript

'use client';
import { useState, useEffect } from 'react';
import { useRouter, useSearchParams } from 'next/navigation';
import { toast } from 'react-hot-toast';
import { Maximize, Download, Sparkles, Trash2, RefreshCw } from 'lucide-react';
import FileUpload from '@/components/FileUpload';
import { modulesApi, assetsApi, jobsApi } from '@/lib/api';
import { useStore } from '@/lib/store';
import { useDragFromCarousel } from '@/hooks/useDragFromCarousel';
const scaleOptions = [
{ value: 2, label: '2x' },
{ value: 4, label: '4x' },
{ value: 6, label: '6x' },
];
const modelOptions = [
{ value: 'Standard V2', label: 'Standard V2' },
{ value: 'High Fidelity V2', label: 'High Fidelity V2' },
{ value: 'Low Resolution V2', label: 'Low Resolution V2' },
{ value: 'CGI', label: 'CGI' },
{ value: 'Text Refine', label: 'Text Refine' },
{ value: 'Enhance Generative', label: 'Enhance Generative' },
{ value: 'Auto', label: 'Auto' },
];
interface QueueItem {
id: string;
file?: File;
assetId?: string;
filename?: string; // fallback for display if file is missing
jobId?: string;
outputAssetId?: string;
status: 'pending' | 'uploading' | 'processing' | 'completed' | 'error';
error?: string;
result?: any;
}
export default function ImageUpscalePage() {
const router = useRouter();
const searchParams = useSearchParams();
const { addJob, updateJob } = useStore();
const [mounted, setMounted] = useState(false);
const [queue, setQueue] = useState<QueueItem[]>([]);
const [processing, setProcessing] = useState(false);
// Settings applied to the batch
const [scale, setScale] = useState(2);
const [model, setModel] = useState('Auto');
const [denoiseStrength, setDenoiseStrength] = useState(0.5);
const [sharpen, setSharpen] = useState(0.5);
useEffect(() => {
setMounted(true);
}, []);
useEffect(() => {
// Check for 'assetIds' (batch) OR 'assetId' (legacy/single) in URL
const assetIdsParam = searchParams.get('assetIds');
const singleAssetId = searchParams.get('assetId'); // Legacy support
if ((assetIdsParam || singleAssetId) && assetsApi) {
const idsToFetch = assetIdsParam ? assetIdsParam.split(',') : [singleAssetId!];
// Fetch details for all IDs
Promise.all(idsToFetch.map(id => assetsApi.get(id)))
.then((responses) => {
setQueue(prev => {
// Avoid duplicates
const existingIds = new Set(prev.map(p => p.assetId));
const newItems = responses
.map((r: any) => r.data)
.filter((asset: any) => !existingIds.has(asset.id))
.map((asset: any) => ({
id: Math.random().toString(36).substring(7),
assetId: asset.id,
status: 'pending' as const,
filename: asset.original_filename || asset.filename
}));
return [...prev, ...newItems];
});
// Clean URL
router.replace('/image/upscale');
if (idsToFetch.length > 0) toast.success(`${idsToFetch.length} assets added from library`);
})
.catch(err => {
console.error("Failed to load assets", err);
toast.error("Failed to load some assets");
});
}
}, [searchParams]);
// Handle drag-and-drop from carousel - MUST be at top level, not conditional
useDragFromCarousel({
onAssetDrop: async (assetIds) => {
try {
const responses = await Promise.all(assetIds.map(id => assetsApi.get(id)));
const newItems = responses
.map((r: any) => r.data)
.filter((asset: any) => {
// Only add image assets
return asset.file_type === 'image';
})
.map((asset: any) => ({
id: Math.random().toString(36).substring(7),
assetId: asset.id,
status: 'pending' as const,
filename: asset.original_filename || asset.filename
}));
if (newItems.length > 0) {
setQueue(prev => [...prev, ...newItems]);
toast.success(`${newItems.length} images added from carousel`);
} else {
toast.error('No valid images in selection');
}
} catch (err) {
console.error('Failed to load dragged assets', err);
toast.error('Failed to add assets');
}
},
enabled: mounted
});
if (!mounted) {
return null;
}
const handleFileUpload = (files: File[]) => {
const newItems: QueueItem[] = files.map(file => ({
id: Math.random().toString(36).substring(7),
file,
status: 'pending'
}));
setQueue(prev => [...prev, ...newItems]);
toast.success(`${files.length} images added to queue`);
};
const processItem = async (item: QueueItem) => {
if (item.status === 'completed' || item.status === 'processing') return item;
try {
// 1. Upload if needed
let assetId = item.assetId;
if (!assetId) {
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'uploading' } : i));
if (!item.file) throw new Error("No file to upload");
try {
const uploadRes = await assetsApi.upload(item.file);
assetId = uploadRes.data.id;
} catch (uploadErr: any) {
if (uploadErr.response?.status === 409) {
const existingAssetId = uploadErr.response.data.detail.asset_id;
const shouldOverwrite = window.confirm(
`File "${item.file?.name}" already exists. \nClick OK to Overwrite, Cancel to Use Existing.`
);
if (shouldOverwrite) {
const uploadRes = await assetsApi.upload(item.file, undefined, false, true); // overwrite=true
assetId = uploadRes.data.id;
} else {
assetId = existingAssetId;
}
} else {
throw uploadErr;
}
}
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, assetId, status: 'pending' } : i));
}
// 2. Start Upscale Job
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'processing' } : i));
const response = await modulesApi.upscaleImage({
asset_id: assetId,
scale,
model,
denoise_strength: denoiseStrength,
sharpen,
});
const job = response.data;
addJob({
id: job.id,
module: 'image_upscaling',
status: job.status,
progress: job.progress,
created_at: job.created_at,
});
// Poll locally for this item
let currentJob = job;
while (currentJob.status !== 'completed' && currentJob.status !== 'failed') {
await new Promise(resolve => setTimeout(resolve, 2000));
const pollRes = await jobsApi.get(job.id);
if (pollRes?.data) currentJob = pollRes.data;
else break;
}
if (currentJob.status === 'completed' && currentJob.output_asset_ids?.[0]) {
// Fetch the output asset to get details
const assetRes = await assetsApi.get(currentJob.output_asset_ids[0]);
const outputAsset = assetRes.data;
setQueue(prev => prev.map(i => i.id === item.id ? {
...i,
status: 'completed',
jobId: job.id,
outputAssetId: outputAsset.id,
result: outputAsset
} : i));
return { ...item, status: 'completed', result: outputAsset };
} else {
throw new Error(currentJob.error_message || 'Job failed');
}
} catch (err: any) {
console.error(err);
const errorMsg = err.message || 'Failed';
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'error', error: errorMsg } : i));
return { ...item, status: 'error', error: errorMsg };
}
};
const handleProcessQueue = async () => {
setProcessing(true);
const pending = queue.filter(i => i.status === 'pending' || i.status === 'error');
// Process strictly sequentially (concurrency 1) to avoid rate limits on Topaz
// Topaz can be touchy with concurrent heavy upscales
const limit = 1;
for (let i = 0; i < pending.length; i += limit) {
const chunk = pending.slice(i, i + limit);
await Promise.all(chunk.map(item => processItem(item)));
}
setProcessing(false);
toast.success('Batch processing complete');
};
const handleClearQueue = () => {
setQueue([]);
};
const removeItem = (id: string) => {
setQueue(prev => prev.filter(i => i.id !== id));
};
const handleDownload = async (item: QueueItem) => {
if (!item.result) return;
try {
// Direct download link logic
const url = `/api/v1/assets/${item.result.id}/download`;
const link = document.createElement('a');
link.href = url;
link.download = item.result.original_filename; // Browser should handle this
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
} catch (err) {
toast.error('Failed to download image');
}
};
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">
<Maximize className="w-6 h-6 text-forge-yellow" />
</div>
<div>
<h1 className="text-2xl font-bold text-white">Image Upscaler</h1>
<p className="text-gray-500">Enhance multiple images with Topaz Labs AI</p>
</div>
</div>
<div className="grid grid-cols-1 lg:grid-cols-3 gap-8">
{/* Left Column: Settings */}
<div className="space-y-6 lg:col-span-1">
<div className="bg-forge-dark p-6 rounded-xl border border-gray-800 space-y-6">
<h3 className="text-lg font-semibold text-white">Batch Settings</h3>
{/* Scale */}
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Scale Factor
</label>
<div className="flex gap-2">
{scaleOptions.map((option) => (
<button
key={option.value}
onClick={() => setScale(option.value)}
className={`flex-1 py-2 rounded-lg font-medium transition-colors text-sm ${scale === option.value
? 'bg-forge-yellow text-black'
: 'bg-black/40 border border-gray-700 text-gray-300 hover:border-gray-600'
}`}
>
{option.label}
</button>
))}
</div>
</div>
{/* Model */}
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Upscaling Model
</label>
<select
value={model}
onChange={(e: React.ChangeEvent<HTMLSelectElement>) => setModel(e.target.value)}
className="select-field w-full"
>
{modelOptions.map((option) => (
<option key={option.value} value={option.value}>
{option.label}
</option>
))}
</select>
</div>
{/* Denoise */}
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Denoise Strength: {denoiseStrength.toFixed(1)}
</label>
<input
type="range"
min={0}
max={1}
step={0.1}
value={denoiseStrength}
onChange={(e: React.ChangeEvent<HTMLInputElement>) => setDenoiseStrength(parseFloat(e.target.value))}
className="w-full accent-forge-yellow"
/>
</div>
{/* Sharpen */}
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Sharpen: {sharpen.toFixed(1)}
</label>
<input
type="range"
min={0}
max={1}
step={0.1}
value={sharpen}
onChange={(e: React.ChangeEvent<HTMLInputElement>) => setSharpen(parseFloat(e.target.value))}
className="w-full accent-forge-yellow"
/>
</div>
</div>
</div>
{/* Right Column: Queue & Upload */}
<div className="lg:col-span-2 space-y-6">
{/* Upload Area */}
<div>
<FileUpload
onUploadMultiple={handleFileUpload}
accept={{ 'image/*': ['.png', '.jpg', '.jpeg', '.webp'] }}
label="Drag & drop images here (Multiple allowed)"
multiple={true}
/>
</div>
{/* Queue Actions */}
{queue.length > 0 && (
<div className="flex flex-wrap gap-4 items-center justify-between bg-forge-dark p-4 rounded-xl border border-gray-800">
<div className="text-white font-medium">
Queue: {queue.length} items ({queue.filter(i => i.status === 'completed').length} completed)
</div>
<div className="flex gap-2">
<button
onClick={handleClearQueue}
className="btn-secondary text-sm flex items-center gap-2"
disabled={processing}
>
<Trash2 className="w-4 h-4" /> Clear
</button>
<button
onClick={handleProcessQueue}
disabled={processing || !queue.some(i => i.status === 'pending' || i.status === 'error')}
className="btn-primary text-sm flex items-center gap-2 disabled:opacity-50"
>
{processing ? <RefreshCw className="w-4 h-4 animate-spin" /> : <Sparkles className="w-4 h-4" />}
{processing ? 'Processing...' : 'Process Batch'}
</button>
</div>
</div>
)}
{/* Queue List */}
<div className="space-y-4">
{queue.map(item => (
<div key={item.id} className="bg-forge-dark rounded-xl border border-gray-800 p-4 flex gap-4 items-center">
{/* Thumbnail */}
<div className="w-20 h-20 bg-black/40 rounded-lg flex-shrink-0 overflow-hidden relative border border-gray-800">
<img
src={item.file ? URL.createObjectURL(item.file) : `/api/v1/assets/${item.assetId}/download`}
alt="thumb"
className="w-full h-full object-cover"
/>
</div>
{/* Info */}
<div className="flex-1 min-w-0">
<h4 className="text-white font-medium truncate">{item.file?.name || item.filename || 'Unknown File'}</h4>
<div className="text-sm mt-1">
{item.status === 'pending' && <span className="text-gray-500">Waiting to start...</span>}
{item.status === 'uploading' && <span className="text-blue-400">Uploading original...</span>}
{item.status === 'processing' && <span className="text-forge-yellow flex items-center gap-2"><RefreshCw className="w-3 h-3 animate-spin" /> Upscaling...</span>}
{item.status === 'completed' && <span className="text-green-400 flex items-center gap-2"><Sparkles className="w-3 h-3" /> Complete</span>}
{item.status === 'error' && <span className="text-red-400">Error: {item.error}</span>}
</div>
</div>
{/* Actions */}
<div className="flex items-center gap-2">
{item.status === 'completed' && item.result && (
<button
onClick={() => handleDownload(item)}
className="p-2 hover:bg-white/10 rounded text-forge-yellow transition-colors"
title="Download"
>
<Download className="w-5 h-5" />
</button>
)}
<button
onClick={() => removeItem(item.id)}
className="p-2 hover:bg-white/10 rounded text-gray-500 hover:text-red-400 transition-colors"
disabled={processing && item.status === 'processing'}
>
<Trash2 className="w-5 h-5" />
</button>
</div>
</div>
))}
{queue.length === 0 && !processing && (
<div className="text-center py-12 text-gray-500">
Add images to start batch upscaling
</div>
)}
</div>
</div>
</div>
</div>
);
}