import os import shutil # Define the directory structure and file contents files = { # Backend files "backend/__init__.py": "", "backend/main.py": """import logging from contextlib import asynccontextmanager from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware import uvicorn import asyncio from datetime import datetime from api.routes import router as api_router from services.websocket_manager import WebSocketManager from services.cache_manager import CacheManager from config.settings import settings from models.model_manager import ModelManager from models.database import init_db # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) websocket_manager = WebSocketManager() cache_manager = CacheManager() model_manager = ModelManager() @asynccontextmanager async def lifespan(app: FastAPI): logger.info("Starting AI Creative Suite...") await init_db() await model_manager.initialize() await cache_manager.initialize() asyncio.create_task(cleanup_old_jobs()) yield await model_manager.cleanup() await cache_manager.cleanup() app = FastAPI(title="AI Creative Suite API", version="1.0.0", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=settings.ALLOWED_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) app.include_router(api_router, prefix="/api/v1") @app.get("/") async def root(): return {"name": "AI Creative Suite API", "status": "operational"} @app.get("/health") async def health_check(): return {"status": "healthy", "models_loaded": model_manager.is_initialized} @app.websocket("/ws/{job_id}") async def websocket_endpoint(websocket: WebSocket, job_id: str): await websocket_manager.connect(websocket, job_id) try: while True: data = await websocket.receive_text() await websocket_manager.handle_message(job_id, data) except WebSocketDisconnect: websocket_manager.disconnect(job_id) async def cleanup_old_jobs(): while True: await asyncio.sleep(3600) await cache_manager.cleanup_expired_jobs() if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=settings.DEBUG) """, "backend/requirements.txt": """fastapi==0.104.1 uvicorn[standard]==0.24.0 pydantic==2.4.2 pydantic-settings==2.0.3 sqlalchemy==2.0.23 psycopg2-binary==2.9.9 redis==5.0.1 celery==5.3.4 boto3==1.28.64 python-multipart==0.0.6 python-jose[cryptography]==3.3.0 passlib[bcrypt]==1.7.4 python-dotenv==1.0.0 alembic==1.12.1 httpx==0.25.1 torch==2.1.0 torchvision==0.16.0 diffusers==0.24.0 transformers==4.35.0 accelerate==0.24.1 xformers==0.0.22 opencv-python==4.8.1.78 pillow==10.1.0 numpy==1.24.3 decord==0.6.0 imageio[ffmpeg]==2.31.6 moviepy==1.0.3 prometheus-client==0.19.0 opentelemetry-api==1.21.0 opentelemetry-sdk==1.21.0 pytest==7.4.3 pytest-asyncio==0.21.1 """, "backend/config/__init__.py": "", "backend/config/settings.py": """from pydantic_settings import BaseSettings from typing import List import os class Settings(BaseSettings): API_VERSION: str = "v1" DEBUG: bool = os.getenv("DEBUG", "False").lower() == "true" ALLOWED_ORIGINS: List[str] = ["http://localhost:3000", "http://localhost:5173"] DATABASE_URL: str = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/ai_creative") REDIS_URL: str = os.getenv("REDIS_URL", "redis://localhost:6379/0") AWS_ACCESS_KEY_ID: str = os.getenv("AWS_ACCESS_KEY_ID", "") AWS_SECRET_ACCESS_KEY: str = os.getenv("AWS_SECRET_ACCESS_KEY", "") S3_BUCKET: str = os.getenv("S3_BUCKET", "ai-creative-outputs") S3_REGION: str = os.getenv("S3_REGION", "us-east-1") MODEL_CACHE_DIR: str = os.getenv("MODEL_CACHE_DIR", "/models") IMAGE_MODEL: str = "stabilityai/stable-diffusion-xl-base-1.0" VIDEO_MODEL: str = "stabilityai/stable-video-diffusion-img2vid" CUDA_VISIBLE_DEVICES: str = os.getenv("CUDA_VISIBLE_DEVICES", "0") TORCH_DTYPE: str = "float16" CELERY_BROKER_URL: str = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/1") CELERY_RESULT_BACKEND: str = os.getenv("CELERY_RESULT_BACKEND", "redis://localhost:6379/2") SECRET_KEY: str = os.getenv("SECRET_KEY", "your-secret-key-change-in-production") ALGORITHM: str = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES: int = 30 class Config: env_file = ".env" settings = Settings() """, "backend/api/__init__.py": "", "backend/api/schemas.py": """from pydantic import BaseModel, Field, validator from typing import Optional, List from datetime import datetime from enum import Enum class JobType(str, Enum): IMAGE = "image" VIDEO = "video" class StyleType(str, Enum): REALISTIC = "realistic" CINEMATIC = "cinematic" ANIME = "anime" THREE_D = "3d" PAINTING = "painting" CONCEPT_ART = "concept" class GenerationRequest(BaseModel): prompt: str = Field(..., min_length=1, max_length=1000) negative_prompt: Optional[str] = Field("", max_length=500) style: StyleType = StyleType.REALISTIC resolution: str = "1024x1024" num_outputs: int = Field(1, ge=1, le=4) seed: Optional[int] = Field(None, ge=0, le=2**32-1) guidance_scale: float = Field(7.5, ge=1, le=15) num_inference_steps: int = Field(50, ge=10, le=100) @validator('resolution') def validate_resolution(cls, v): valid = ['512x512', '768x768', '1024x1024', '1024x768', '768x1024'] if v not in valid: raise ValueError(f'Resolution must be one of {valid}') return v class VideoGenerationRequest(GenerationRequest): duration_seconds: int = Field(5, ge=1, le=10) fps: int = Field(30, ge=24, le=60) motion_intensity: float = Field(0.5, ge=0, le=1) camera_angle: str = "dynamic" reference_image: Optional[str] = None class GenerationResponse(BaseModel): job_id: str status: str output_urls: Optional[List[str]] = None created_at: datetime estimated_completion: Optional[datetime] = None class JobStatusResponse(BaseModel): job_id: str status: str progress: int = 0 output_urls: Optional[List[str]] = None error_message: Optional[str] = None created_at: datetime completed_at: Optional[datetime] = None class UserCreate(BaseModel): email: str username: str password: str class UserResponse(BaseModel): id: str email: str username: str credits_balance: int created_at: datetime class Token(BaseModel): access_token: str token_type: str """, "backend/api/routes.py": """from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks from typing import List import uuid from datetime import datetime, timedelta from .schemas import ( GenerationRequest, GenerationResponse, VideoGenerationRequest, JobStatusResponse, UserCreate, UserResponse, Token ) from services.auth import get_current_user, create_access_token, verify_password, get_password_hash from services.task_manager import TaskManager from models.database import get_db, User, GenerationJob from workers.tasks import generate_image_task, generate_video_task from sqlalchemy.orm import Session from config.settings import settings router = APIRouter() task_manager = TaskManager() @router.post("/auth/register", response_model=UserResponse) async def register(user_data: UserCreate, db: Session = Depends(get_db)): existing = db.query(User).filter((User.email == user_data.email) | (User.username == user_data.username)).first() if existing: raise HTTPException(400, "User already exists") user = User( id=str(uuid.uuid4()), email=user_data.email, username=user_data.username, hashed_password=get_password_hash(user_data.password), credits_balance=100 ) db.add(user) db.commit() db.refresh(user) return user @router.post("/auth/login", response_model=Token) async def login(username: str, password: str, db: Session = Depends(get_db)): user = db.query(User).filter(User.username == username).first() if not user or not verify_password(password, user.hashed_password): raise HTTPException(401, "Invalid credentials") access_token = create_access_token(data={"sub": user.username}) return {"access_token": access_token, "token_type": "bearer"} @router.post("/generate/image", response_model=GenerationResponse) async def generate_image( request: GenerationRequest, background_tasks: BackgroundTasks, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): if current_user.credits_balance < request.num_outputs: raise HTTPException(402, "Insufficient credits") job_id = str(uuid.uuid4()) job = GenerationJob( id=job_id, user_id=current_user.id, job_type="image", status="queued", prompt=request.prompt, negative_prompt=request.negative_prompt, parameters=request.dict(), created_at=datetime.utcnow() ) db.add(job) db.commit() current_user.credits_balance -= request.num_outputs db.commit() generate_image_task.delay(job_id=job_id, request=request.dict(), user_id=current_user.id) return GenerationResponse( job_id=job_id, status="queued", created_at=job.created_at, estimated_completion=datetime.utcnow() + timedelta(seconds=30) ) @router.post("/generate/video", response_model=GenerationResponse) async def generate_video( request: VideoGenerationRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): credit_cost = request.duration_seconds * 2 if current_user.credits_balance < credit_cost: raise HTTPException(402, "Insufficient credits") job_id = str(uuid.uuid4()) job = GenerationJob( id=job_id, user_id=current_user.id, job_type="video", status="queued", prompt=request.prompt, negative_prompt=request.negative_prompt, parameters=request.dict(), created_at=datetime.utcnow() ) db.add(job) db.commit() current_user.credits_balance -= credit_cost db.commit() generate_video_task.delay(job_id=job_id, request=request.dict(), user_id=current_user.id) return GenerationResponse( job_id=job_id, status="queued", created_at=job.created_at, estimated_completion=datetime.utcnow() + timedelta(minutes=2) ) @router.get("/job/{job_id}", response_model=JobStatusResponse) async def get_job_status(job_id: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first() if not job or job.user_id != current_user.id: raise HTTPException(404, "Job not found") return JobStatusResponse( job_id=job.id, status=job.status, progress=getattr(job, 'progress', 0), output_urls=job.output_urls, error_message=job.error_message, created_at=job.created_at, completed_at=job.completed_at ) @router.get("/user/credits") async def get_credits(current_user: User = Depends(get_current_user)): return {"credits": current_user.credits_balance} """, "backend/models/__init__.py": "", "backend/models/database.py": """from sqlalchemy import create_engine, Column, String, DateTime, Integer, Float, JSON, Text from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker from datetime import datetime import uuid from config.settings import settings Base = declarative_base() class GenerationJob(Base): __tablename__ = "generation_jobs" id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) user_id = Column(String(36), nullable=False, index=True) job_type = Column(String(50), nullable=False) status = Column(String(50), default="queued") prompt = Column(Text, nullable=False) negative_prompt = Column(Text) parameters = Column(JSON) output_urls = Column(JSON) error_message = Column(Text) progress = Column(Integer, default=0) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) completed_at = Column(DateTime) class User(Base): __tablename__ = "users" id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) email = Column(String(255), unique=True, nullable=False, index=True) username = Column(String(255), unique=True, nullable=False, index=True) hashed_password = Column(String(255), nullable=False) credits_balance = Column(Integer, default=100) created_at = Column(DateTime, default=datetime.utcnow) last_login = Column(DateTime) engine = create_engine(settings.DATABASE_URL) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) async def init_db(): Base.metadata.create_all(bind=engine) def get_db(): db = SessionLocal() try: yield db finally: db.close() """, "backend/models/model_manager.py": """import torch import asyncio from typing import Optional, List from diffusers import StableDiffusionXLPipeline, StableVideoDiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL from PIL import Image import logging import gc logger = logging.getLogger(__name__) class ModelManager: def __init__(self): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.image_model = None self.video_model = None self.is_initialized = False self.lock = asyncio.Lock() async def initialize(self): async with self.lock: if not self.is_initialized: logger.info(f"Initializing models on {self.device}") await self._init_image_model() if torch.cuda.is_available(): await self._init_video_model() self.is_initialized = True async def _init_image_model(self): try: vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) self.image_model = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True ) self.image_model.to(self.device) if torch.cuda.is_available(): self.image_model.enable_model_cpu_offload() self.image_model.enable_vae_slicing() self.image_model.enable_vae_tiling() self.image_model.enable_attention_slicing() self.image_model.scheduler = DPMSolverMultistepScheduler.from_config( self.image_model.scheduler.config, use_karras_sigmas=True ) logger.info("Image model loaded") except Exception as e: logger.error(f"Failed to load image model: {e}") async def _init_video_model(self): try: self.video_model = StableVideoDiffusionPipeline.from_pretrained( "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16" ) self.video_model.to(self.device) self.video_model.enable_model_cpu_offload() logger.info("Video model loaded") except Exception as e: logger.error(f"Failed to load video model: {e}") async def generate_image(self, prompt, negative_prompt="", width=1024, height=1024, num_inference_steps=50, guidance_scale=7.5, num_images=1, seed=None, style="realistic"): if not self.is_initialized: await self.initialize() enhanced_prompt = self._enhance_prompt(prompt, style) generator = torch.Generator(device=self.device).manual_seed(seed) if seed else None with torch.no_grad(), torch.autocast("cuda"): result = self.image_model( prompt=enhanced_prompt, negative_prompt=negative_prompt, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_images, generator=generator ) return result.images def _enhance_prompt(self, prompt, style): style_map = { "realistic": "photorealistic, high quality, detailed", "cinematic": "cinematic, dramatic lighting, shallow depth of field", "anime": "anime style, vibrant colors, cel shading", "3d": "3d render, blender, octane render", "painting": "oil painting, brush strokes, masterpiece" } style_text = style_map.get(style, style_map["realistic"]) return f"{prompt}, {style_text}, 8k, masterpiece" async def cleanup(self): if self.image_model: del self.image_model if self.video_model: del self.video_model if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() self.is_initialized = False """, "backend/services/__init__.py": "", "backend/services/auth.py": """from datetime import datetime, timedelta from jose import JWTError, jwt from passlib.context import CryptContext from fastapi import Depends, HTTPException, status from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from sqlalchemy.orm import Session from models.database import get_db, User from config.settings import settings pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") security = HTTPBearer() def verify_password(plain, hashed): return pwd_context.verify(plain, hashed) def get_password_hash(password): return pwd_context.hash(password) def create_access_token(data: dict, expires_delta: timedelta = None): to_encode = data.copy() expire = datetime.utcnow() + (expires_delta or timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)) to_encode.update({"exp": expire}) return jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM) def decode_access_token(token: str): try: return jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM]) except JWTError: return None async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security), db: Session = Depends(get_db)): token = credentials.credentials payload = decode_access_token(token) if not payload: raise HTTPException(status_code=401, detail="Invalid token") username = payload.get("sub") user = db.query(User).filter(User.username == username).first() if not user: raise HTTPException(status_code=401, detail="User not found") return user """, "backend/services/storage.py": """import boto3 from PIL import Image import io import logging from datetime import datetime from config.settings import settings logger = logging.getLogger(__name__) class StorageService: def __init__(self): self.s3_client = None if settings.AWS_ACCESS_KEY_ID: self.s3_client = boto3.client( 's3', aws_access_key_id=settings.AWS_ACCESS_KEY_ID, aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY, region_name=settings.S3_REGION ) async def upload_image(self, image: Image.Image, job_id: str, index: int) -> str: buffer = io.BytesIO() image.save(buffer, format='PNG', optimize=True) buffer.seek(0) timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") key = f"images/{job_id}/{timestamp}_{index}.png" if self.s3_client: self.s3_client.put_object(Bucket=settings.S3_BUCKET, Key=key, Body=buffer.getvalue(), ContentType='image/png') return f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/{key}" else: import os os.makedirs(f"/tmp/images/{job_id}", exist_ok=True) filepath = f"/tmp/images/{job_id}/{timestamp}_{index}.png" image.save(filepath) return f"file://{filepath}" async def upload_video(self, frames, job_id: str, fps: int) -> str: # Simplified: just return placeholder return f"file:///tmp/videos/{job_id}/video.mp4" """, "backend/services/websocket_manager.py": """from fastapi import WebSocket from typing import Dict, Set import asyncio import json import logging logger = logging.getLogger(__name__) class WebSocketManager: def __init__(self): self.active_connections: Dict[str, Set[WebSocket]] = {} self.lock = asyncio.Lock() async def connect(self, websocket: WebSocket, job_id: str): await websocket.accept() async with self.lock: if job_id not in self.active_connections: self.active_connections[job_id] = set() self.active_connections[job_id].add(websocket) logger.info(f"WebSocket connected for job {job_id}") def disconnect(self, job_id: str): # handled elsewhere pass async def send_progress(self, job_id: str, progress: int, message: str = None): data = {"type": "progress", "progress": progress, "message": message} await self._send_to_job(job_id, data) async def send_complete(self, job_id: str, urls: list): data = {"type": "complete", "urls": urls} await self._send_to_job(job_id, data) async def send_error(self, job_id: str, error: str): data = {"type": "error", "error": error} await self._send_to_job(job_id, data) async def _send_to_job(self, job_id: str, data: dict): async with self.lock: if job_id not in self.active_connections: return connections = self.active_connections[job_id].copy() for conn in connections: try: await conn.send_json(data) except: pass async def handle_message(self, job_id: str, message: str): try: data = json.loads(message) if data.get("type") == "ping": await self._send_to_job(job_id, {"type": "pong"}) except: pass """, "backend/services/task_manager.py": """class TaskManager: async def create_job(self, job, db): # Placeholder pass async def get_job(self, job_id, db): # Placeholder pass async def create_project(self, user_id, name, description, db): # Placeholder return "project-id" """, "backend/workers/__init__.py": "", "backend/workers/tasks.py": """from celery import Celery import asyncio import logging from datetime import datetime from config.settings import settings from models.model_manager import ModelManager from models.database import SessionLocal, GenerationJob from services.storage import StorageService app = Celery('tasks', broker=settings.CELERY_BROKER_URL) app.conf.update(task_serializer='json', result_serializer='json', accept_content=['json']) logger = logging.getLogger(__name__) model_manager = ModelManager() storage_service = StorageService() @app.task(bind=True) def generate_image_task(self, job_id: str, request: dict, user_id: str): db = SessionLocal() try: job = db.query(GenerationJob).filter(GenerationJob.id == job_id).first() if job: job.status = "processing" db.commit() loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) images = loop.run_until_complete(model_manager.generate_image(**request)) urls = [] for i, img in enumerate(images): url = loop.run_until_complete(storage_service.upload_image(img, job_id, i)) urls.append(url) if job: job.status = "completed" job.output_urls = urls job.completed_at = datetime.utcnow() db.commit() return {'urls': urls} except Exception as e: logger.error(f"Task failed: {e}") if job: job.status = "failed" job.error_message = str(e) db.commit() raise finally: db.close() """, "frontend/package.json": """{ "name": "ai-creative-suite-frontend", "version": "1.0.0", "private": true, "scripts": { "dev": "next dev", "build": "next build", "start": "next start" }, "dependencies": { "react": "^18.2.0", "react-dom": "^18.2.0", "next": "14.0.0", "@tanstack/react-query": "^5.8.4", "axios": "^1.6.2", "socket.io-client": "^4.5.4", "framer-motion": "^10.16.5", "tailwindcss": "^3.3.6", "react-hot-toast": "^2.4.1", "lucide-react": "^0.294.0" }, "devDependencies": { "@types/react": "^18.2.42", "@types/node": "^20.10.4", "typescript": "^5.3.2", "eslint": "^8.55.0", "eslint-config-next": "14.0.0", "autoprefixer": "^10.4.16", "postcss": "^8.4.32" } } """, "frontend/next.config.js": """/** @type {import('next').NextConfig} */ const nextConfig = { reactStrictMode: true, swcMinify: true, async rewrites() { return [{ source: '/api/:path*', destination: 'http://localhost:8000/api/:path*' }]; }, }; module.exports = nextConfig; """, "frontend/tailwind.config.js": """/** @type {import('tailwindcss').Config} */ module.exports = { content: ['./src/**/*.{js,ts,jsx,tsx}'], theme: { extend: {} }, plugins: [], }; """, "frontend/tsconfig.json": """{ "compilerOptions": { "target": "es5", "lib": ["dom", "dom.iterable", "esnext"], "allowJs": true, "skipLibCheck": true, "strict": true, "forceConsistentCasingInFileNames": true, "noEmit": true, "esModuleInterop": true, "module": "esnext", "moduleResolution": "node", "resolveJsonModule": true, "isolatedModules": true, "jsx": "preserve", "incremental": true, "paths": { "@/*": ["./src/*"] } }, "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx"], "exclude": ["node_modules"] } """, "frontend/src/pages/_app.tsx": """import type { AppProps } from 'next/app'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { Toaster } from 'react-hot-toast'; import '../styles/globals.css'; const queryClient = new QueryClient(); export default function App({ Component, pageProps }: AppProps) { return ( ); } """, "frontend/src/pages/index.tsx": """import { useState } from 'react'; import GenerationInterface from '../components/GenerationInterface'; import { motion } from 'framer-motion'; export default function Home() { const [activeTab, setActiveTab] = useState<'image' | 'video'>('image'); return (
); } """, "frontend/src/components/GenerationInterface.tsx": """import React, { useState, useRef, useEffect } from 'react'; import { useMutation } from '@tanstack/react-query'; import { motion, AnimatePresence } from 'framer-motion'; import toast from 'react-hot-toast'; import { Sparkles, Download, RefreshCw } from 'lucide-react'; import axios from 'axios'; interface Props { type: 'image' | 'video'; } const GenerationInterface: React.FC = ({ type }) => { const [prompt, setPrompt] = useState(''); const [negativePrompt, setNegativePrompt] = useState(''); const [selectedStyle, setSelectedStyle] = useState('realistic'); const [resolution, setResolution] = useState('1024x1024'); const [numOutputs, setNumOutputs] = useState(1); const [guidanceScale, setGuidanceScale] = useState(7.5); const [duration, setDuration] = useState(5); const [motionIntensity, setMotionIntensity] = useState(0.5); const [currentJob, setCurrentJob] = useState(null); const [generatedUrls, setGeneratedUrls] = useState([]); const [progress, setProgress] = useState(0); const wsRef = useRef(null); const styles = ['realistic', 'cinematic', 'anime', '3d', 'painting', 'concept']; const resolutions = ['512x512', '768x768', '1024x1024', '1024x768', '768x1024']; const generateMutation = useMutation({ mutationFn: async (params: any) => { const endpoint = type === 'image' ? '/api/v1/generate/image' : '/api/v1/generate/video'; const res = await axios.post(endpoint, params, { headers: { Authorization: `Bearer ${localStorage.getItem('token')}` } }); return res.data; }, onSuccess: (data) => { setCurrentJob(data.job_id); connectWebSocket(data.job_id); toast.success('Generation started!'); }, onError: () => toast.error('Failed to start generation'), }); const connectWebSocket = (jobId: string) => { const ws = new WebSocket(`ws://localhost:8000/ws/${jobId}`); ws.onmessage = (event) => { const data = JSON.parse(event.data); if (data.type === 'progress') setProgress(data.progress); else if (data.type === 'complete') { setGeneratedUrls(prev => [...prev, ...data.urls]); setCurrentJob(null); setProgress(0); toast.success('Complete!'); ws.close(); } else if (data.type === 'error') { toast.error(data.error); setCurrentJob(null); ws.close(); } }; wsRef.current = ws; }; const handleGenerate = () => { if (!prompt.trim()) return toast.error('Enter a prompt'); const params: any = { prompt, negative_prompt: negativePrompt, style: selectedStyle, resolution, num_outputs: numOutputs, guidance_scale: guidanceScale }; if (type === 'video') { params.duration_seconds = duration; params.motion_intensity = motionIntensity; params.fps = 30; } generateMutation.mutate(params); }; return (

{type === 'image' ? 'Create Image' : 'Create Video'}