import React, { useState, useEffect, useCallback } from 'react'; import { PolicyBadge } from './PolicyBadge'; import { resolveFileUrl } from '../../resolveFileUrl'; const DEFAULT_SETTINGS = { model: 'abyssOrangeMix3_aom3a1b.safetensors', width: 512, height: 768, steps: 25, cfg: 6.0, sampler: 'dpm++_2m_karras', clipSkip: 2, seed: 'random', }; const MODELS = [ { id: 'abyssOrangeMix3_aom3a1b.safetensors', label: 'AbyssOrangeMix3 (AOM3)', nsfw: true }, { id: 'counterfeit_v30.safetensors', label: 'Counterfeit V3.0', nsfw: true }, { id: 'anything_v5PrtRE.safetensors', label: 'Anything V5', nsfw: true }, { id: 'sd_xl_base_1.0.safetensors', label: 'SDXL Base 1.0', nsfw: false }, { id: 'dreamshaper_8.safetensors', label: 'DreamShaper 8', nsfw: true }, ]; const SIZES = [ { width: 512, height: 512, label: 'Square (512x512)' }, { width: 512, height: 768, label: 'Portrait (512x768)' }, { width: 768, height: 512, label: 'Landscape (768x512)' }, { width: 768, height: 1024, label: 'Tall (768x1024)' }, { width: 1024, height: 768, label: 'Wide (1024x768)' }, ]; const SAMPLERS = [ 'dpm++_2m_karras', 'euler_a', 'euler', 'dpm++_sde_karras', 'ddim', ]; /** * Image Generator Interface * * Main generation UI with prompt input, settings, and preview. * Includes real-time policy checking. */ export const ImageGenerator = ({ contentRating, onGenerate, isGenerating = false, generatedImage = null, presetSettings, presetPositivePrefix = '', presetPositiveSuffix = '', presetNegative = '', }) => { const [prompt, setPrompt] = useState(''); const [negativePrompt, setNegativePrompt] = useState(presetNegative); const [settings, setSettings] = useState({ ...DEFAULT_SETTINGS, ...presetSettings, }); const [policyResult, setPolicyResult] = useState(null); const [showAdvanced, setShowAdvanced] = useState(false); // Update negative prompt when preset changes useEffect(() => { if (presetNegative) { setNegativePrompt(presetNegative); } }, [presetNegative]); // Apply preset settings useEffect(() => { if (presetSettings) { setSettings((prev) => ({ ...prev, ...presetSettings })); } }, [presetSettings]); // Real-time policy check with debounce useEffect(() => { const fullPrompt = `${presetPositivePrefix}${prompt}${presetPositiveSuffix}`; if (!fullPrompt.trim()) { setPolicyResult(null); return; } const timer = setTimeout(async () => { try { const res = await fetch(`/studio/image/policy-check?prompt=${encodeURIComponent(fullPrompt)}&content_rating=${contentRating}&provider=comfyui`, { method: 'POST' }); const data = await res.json(); setPolicyResult(data); } catch (e) { console.error('Policy check failed:', e); } }, 300); return () => clearTimeout(timer); }, [prompt, contentRating, presetPositivePrefix, presetPositiveSuffix]); const handleGenerate = useCallback(() => { const fullPrompt = `${presetPositivePrefix}${prompt}${presetPositiveSuffix}`; onGenerate(fullPrompt, negativePrompt, settings); }, [prompt, negativePrompt, settings, presetPositivePrefix, presetPositiveSuffix, onGenerate]); const handleRandomize = () => { setSettings((prev) => ({ ...prev, seed: Math.floor(Math.random() * 2147483647), })); }; const selectedSize = SIZES.find((s) => s.width === settings.width && s.height === settings.height); const canGenerate = prompt.trim() && (!policyResult || policyResult.allowed); return (
Generating...
Enter a prompt to generate