| import os |
| import shutil |
|
|
| |
| 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 ( |
| <QueryClientProvider client={queryClient}> |
| <Component {...pageProps} /> |
| <Toaster position="top-right" /> |
| </QueryClientProvider> |
| ); |
| } |
| """, |
| "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 ( |
| <div className="min-h-screen bg-gradient-to-br from-gray-900 to-gray-800"> |
| <nav className="bg-black bg-opacity-50 backdrop-blur-lg border-b border-gray-700"> |
| <div className="container mx-auto px-4 py-4 flex justify-between"> |
| <span className="text-xl font-bold text-white">AI Creative Suite</span> |
| <div className="flex space-x-4"> |
| <button className="text-gray-300">Dashboard</button> |
| <button className="text-gray-300">Credits: 100</button> |
| </div> |
| </div> |
| </nav> |
| <div className="container mx-auto px-4 py-8"> |
| <div className="flex space-x-4 mb-8"> |
| <button onClick={() => setActiveTab('image')} className={`px-6 py-2 rounded-lg ${activeTab === 'image' ? 'bg-blue-600' : 'bg-gray-700'}`}>Image</button> |
| <button onClick={() => setActiveTab('video')} className={`px-6 py-2 rounded-lg ${activeTab === 'video' ? 'bg-blue-600' : 'bg-gray-700'}`}>Video</button> |
| </div> |
| <motion.div key={activeTab} initial={{ opacity: 0, y: 20 }} animate={{ opacity: 1, y: 0 }}> |
| <GenerationInterface type={activeTab} /> |
| </motion.div> |
| </div> |
| </div> |
| ); |
| } |
| """, |
| "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<Props> = ({ 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<string | null>(null); |
| const [generatedUrls, setGeneratedUrls] = useState<string[]>([]); |
| const [progress, setProgress] = useState(0); |
| const wsRef = useRef<WebSocket | null>(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 ( |
| <div className="grid grid-cols-1 lg:grid-cols-2 gap-8"> |
| <div className="bg-gray-800 rounded-xl p-6 shadow-xl"> |
| <h2 className="text-2xl font-bold text-white mb-6">{type === 'image' ? 'Create Image' : 'Create Video'}</h2> |
| <div className="mb-4"><label className="text-sm text-gray-300">Prompt</label><textarea value={prompt} onChange={e => setPrompt(e.target.value)} className="w-full bg-gray-700 rounded-lg p-3 text-white" rows={3} placeholder="Describe..." /></div> |
| <div className="mb-4"><label className="text-sm text-gray-300">Negative Prompt</label><input type="text" value={negativePrompt} onChange={e => setNegativePrompt(e.target.value)} className="w-full bg-gray-700 rounded-lg p-3 text-white" placeholder="Avoid..." /></div> |
| <div className="mb-4"><label className="text-sm text-gray-300">Style</label><div className="grid grid-cols-3 gap-2">{styles.map(s => <button key={s} onClick={() => setSelectedStyle(s)} className={`px-3 py-2 rounded-lg capitalize ${selectedStyle === s ? 'bg-blue-600' : 'bg-gray-700'}`}>{s}</button>)}</div></div> |
| <div className="mb-4"><label className="text-sm text-gray-300">Resolution</label><select value={resolution} onChange={e => setResolution(e.target.value)} className="w-full bg-gray-700 rounded-lg p-3 text-white">{resolutions.map(r => <option key={r}>{r}</option>)}</select></div> |
| <div className="mb-4"><label className="text-sm text-gray-300">Number of Outputs: {numOutputs}</label><input type="range" min={1} max={4} value={numOutputs} onChange={e => setNumOutputs(parseInt(e.target.value))} className="w-full" /></div> |
| <div className="mb-4"><label className="text-sm text-gray-300">Guidance Scale: {guidanceScale}</label><input type="range" min={1} max={15} step={0.5} value={guidanceScale} onChange={e => setGuidanceScale(parseFloat(e.target.value))} className="w-full" /></div> |
| {type === 'video' && (<><div className="mb-4"><label>Duration: {duration}s</label><input type="range" min={1} max={10} value={duration} onChange={e => setDuration(parseInt(e.target.value))} className="w-full" /></div><div className="mb-6"><label>Motion: {motionIntensity}</label><input type="range" min={0} max={1} step={0.05} value={motionIntensity} onChange={e => setMotionIntensity(parseFloat(e.target.value))} className="w-full" /></div></>)} |
| <button onClick={handleGenerate} disabled={generateMutation.isPending || currentJob !== null} className="w-full bg-gradient-to-r from-blue-600 to-purple-600 hover:from-blue-700 hover:to-purple-700 text-white font-bold py-3 px-6 rounded-lg transition disabled:opacity-50 flex items-center justify-center space-x-2"> |
| {generateMutation.isPending || currentJob ? <><RefreshCw className="w-5 h-5 animate-spin" /><span>Generating... {progress}%</span></> : <><Sparkles className="w-5 h-5" /><span>Generate</span></>} |
| </button> |
| {currentJob && <div className="mt-4"><div className="w-full bg-gray-700 rounded-full h-2"><div className="bg-blue-600 h-2 rounded-full transition-all" style={{ width: `${progress}%` }} /></div></div>} |
| </div> |
| <div className="bg-gray-800 rounded-xl p-6 shadow-xl"> |
| <h3 className="text-xl font-bold text-white mb-4">Results</h3> |
| {generatedUrls.length === 0 ? <div className="text-gray-400 text-center py-20">Your generated {type}s will appear here</div> : <div className="grid grid-cols-2 gap-4">{generatedUrls.map((url, i) => (<div key={i} className="relative group"><img src={url} alt="result" className="rounded-lg w-full cursor-pointer" onClick={() => window.open(url)} /><div className="absolute inset-0 bg-black bg-opacity-0 group-hover:bg-opacity-50 transition rounded-lg flex items-center justify-center opacity-0 group-hover:opacity-100"><button onClick={() => window.open(url)} className="bg-white text-gray-900 px-3 py-2 rounded-lg"><Download className="w-4 h-4" /></button></div></div>))}</div>} |
| </div> |
| </div> |
| ); |
| }; |
| export default GenerationInterface; |
| """, |
| "frontend/src/styles/globals.css": """@tailwind base; |
| @tailwind components; |
| @tailwind utilities; |
| body { @apply bg-gray-900 text-white; } |
| ::-webkit-scrollbar { width: 8px; } |
| ::-webkit-scrollbar-track { @apply bg-gray-800; } |
| ::-webkit-scrollbar-thumb { @apply bg-gray-600 rounded-full; } |
| """, |
| "docker-compose.yml": """version: '3.8' |
| services: |
| postgres: |
| image: postgres:15 |
| environment: |
| POSTGRES_USER: ai_user |
| POSTGRES_PASSWORD: ai_pass |
| POSTGRES_DB: ai_creative |
| ports: |
| - "5432:5432" |
| redis: |
| image: redis:7-alpine |
| ports: |
| - "6379:6379" |
| backend: |
| build: ./backend |
| ports: |
| - "8000:8000" |
| environment: |
| DATABASE_URL: postgresql://ai_user:ai_pass@postgres:5432/ai_creative |
| REDIS_URL: redis://redis:6379/0 |
| CELERY_BROKER_URL: redis://redis:6379/1 |
| depends_on: |
| - postgres |
| - redis |
| celery-worker: |
| build: ./backend |
| command: celery -A workers.tasks worker --loglevel=info |
| environment: |
| DATABASE_URL: postgresql://ai_user:ai_pass@postgres:5432/ai_creative |
| REDIS_URL: redis://redis:6379/0 |
| CELERY_BROKER_URL: redis://redis:6379/1 |
| depends_on: |
| - redis |
| - postgres |
| frontend: |
| build: ./frontend |
| ports: |
| - "3000:3000" |
| environment: |
| NEXT_PUBLIC_API_URL: http://localhost:8000 |
| depends_on: |
| - backend |
| """, |
| ".env.example": """DATABASE_URL=postgresql://ai_user:ai_pass@localhost:5432/ai_creative |
| REDIS_URL=redis://localhost:6379/0 |
| SECRET_KEY=your-secret-key-change-in-production |
| DEBUG=True |
| """, |
| "README.md": """# AI Creative Suite |
| |
| ## Setup |
| 1. Install Docker and Docker Compose |
| 2. Copy .env.example to .env |
| 3. Run `docker-compose up -d` |
| 4. Access frontend at http://localhost:3000 |
| 5. API docs at http://localhost:8000/api/docs |
| |
| ## Manual Setup |
| 1. Create virtual env: `python -m venv venv` |
| 2. Activate: `source venv/bin/activate` |
| 3. Install backend: `pip install -r backend/requirements.txt` |
| 4. Run backend: `uvicorn backend.main:app --reload` |
| 5. Run frontend: `cd frontend && npm install && npm run dev` |
| |
| ## Features |
| - Text-to-image generation |
| - Text-to-video generation |
| - User authentication |
| - Real-time progress via WebSockets |
| - Style customization |
| """ |
| } |
|
|
| |
| for filepath, content in files.items(): |
| dirname = os.path.dirname(filepath) |
| if dirname: |
| os.makedirs(dirname, exist_ok=True) |
| with open(filepath, 'w', encoding='utf-8') as f: |
| f.write(content) |
| print(f"Created {filepath}") |
|
|
