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 (