260 lines
10 KiB
Python
260 lines
10 KiB
Python
"""YouTube Creator scene image generation handlers."""
|
|
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional
|
|
import uuid
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
|
|
from middleware.auth_middleware import get_current_user
|
|
from services.database import get_db
|
|
from services.subscription import PricingService
|
|
from services.subscription.preflight_validator import validate_image_generation_operations
|
|
from services.llm_providers.main_image_generation import generate_image
|
|
from services.wavespeed.client import WaveSpeedClient
|
|
from utils.asset_tracker import save_asset_to_library
|
|
from utils.logger_utils import get_service_logger
|
|
|
|
router = APIRouter(tags=["youtube-image"])
|
|
logger = get_service_logger("api.youtube.image")
|
|
|
|
# Directories
|
|
base_dir = Path(__file__).parent.parent.parent.parent
|
|
YOUTUBE_IMAGES_DIR = base_dir / "youtube_images"
|
|
YOUTUBE_IMAGES_DIR.mkdir(parents=True, exist_ok=True)
|
|
YOUTUBE_AVATARS_DIR = base_dir / "youtube_avatars"
|
|
|
|
|
|
class YouTubeImageRequest(BaseModel):
|
|
scene_id: str
|
|
scene_title: Optional[str] = None
|
|
scene_content: Optional[str] = None
|
|
base_avatar_url: Optional[str] = None
|
|
idea: Optional[str] = None
|
|
width: Optional[int] = 1024
|
|
height: Optional[int] = 1024
|
|
custom_prompt: Optional[str] = None
|
|
style: Optional[str] = None # e.g., "Realistic", "Fiction"
|
|
rendering_speed: Optional[str] = None # e.g., "Quality", "Turbo"
|
|
aspect_ratio: Optional[str] = None # e.g., "16:9"
|
|
|
|
|
|
def require_authenticated_user(current_user: Dict[str, Any]) -> str:
|
|
"""Extract and validate user ID from current user."""
|
|
user_id = current_user.get("id") if current_user else None
|
|
if not user_id:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
return str(user_id)
|
|
|
|
|
|
def _load_base_avatar_bytes(avatar_url: str) -> bytes:
|
|
"""Load base avatar bytes for character consistency."""
|
|
filename = avatar_url.split("/")[-1].split("?")[0]
|
|
avatar_path = YOUTUBE_AVATARS_DIR / filename
|
|
if not avatar_path.exists() or not avatar_path.is_file():
|
|
raise HTTPException(status_code=404, detail="Base avatar image not found")
|
|
return avatar_path.read_bytes()
|
|
|
|
|
|
def _save_scene_image(image_bytes: bytes, scene_id: str) -> Dict[str, str]:
|
|
"""Persist generated scene image and return file/url info."""
|
|
unique_id = str(uuid.uuid4())[:8]
|
|
image_filename = f"yt_scene_{scene_id}_{unique_id}.png"
|
|
image_path = YOUTUBE_IMAGES_DIR / image_filename
|
|
with open(image_path, "wb") as f:
|
|
f.write(image_bytes)
|
|
|
|
image_url = f"/api/youtube/images/scenes/{image_filename}"
|
|
return {
|
|
"image_filename": image_filename,
|
|
"image_path": str(image_path),
|
|
"image_url": image_url,
|
|
}
|
|
|
|
|
|
@router.post("/image")
|
|
async def generate_youtube_scene_image(
|
|
request: YouTubeImageRequest,
|
|
current_user: Dict[str, Any] = Depends(get_current_user),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Generate a YouTube scene image, with optional avatar consistency."""
|
|
user_id = require_authenticated_user(current_user)
|
|
|
|
if not request.scene_title:
|
|
raise HTTPException(status_code=400, detail="Scene title is required")
|
|
|
|
try:
|
|
# Pre-flight subscription validation
|
|
pricing_service = PricingService(db)
|
|
validate_image_generation_operations(
|
|
pricing_service=pricing_service,
|
|
user_id=user_id,
|
|
num_images=1,
|
|
)
|
|
logger.info(f"[YouTube] ✅ Pre-flight validation passed for user {user_id}")
|
|
|
|
base_avatar_bytes = None
|
|
if request.base_avatar_url:
|
|
try:
|
|
base_avatar_bytes = _load_base_avatar_bytes(request.base_avatar_url)
|
|
logger.info(f"[YouTube] Loaded base avatar for scene {request.scene_id}")
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"[YouTube] Failed to load base avatar: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "Failed to load base avatar",
|
|
"message": f"Could not load the base avatar image: {str(e)}",
|
|
},
|
|
)
|
|
|
|
# Build prompt
|
|
image_prompt = ""
|
|
if base_avatar_bytes:
|
|
prompt_parts = []
|
|
if request.scene_title:
|
|
prompt_parts.append(f"Scene: {request.scene_title}")
|
|
if request.scene_content:
|
|
content_preview = request.scene_content[:200].replace("\n", " ").strip()
|
|
prompt_parts.append(f"Context: {content_preview}")
|
|
if request.idea:
|
|
prompt_parts.append(f"Video idea: {request.idea[:80].strip()}")
|
|
prompt_parts.append("YouTube creator on camera, engaging and dynamic framing")
|
|
prompt_parts.append("Clean background, good lighting, thumbnail-friendly composition")
|
|
image_prompt = ", ".join(prompt_parts)
|
|
else:
|
|
prompt_parts = [
|
|
"YouTube creator scene",
|
|
"clean, modern background",
|
|
"good lighting, high contrast for thumbnail clarity",
|
|
]
|
|
if request.scene_title:
|
|
prompt_parts.append(f"Scene theme: {request.scene_title}")
|
|
if request.scene_content:
|
|
prompt_parts.append(f"Context: {request.scene_content[:120].replace(chr(10), ' ')}")
|
|
if request.idea:
|
|
prompt_parts.append(f"Topic: {request.idea[:80]}")
|
|
prompt_parts.append("video-optimized composition, 16:9 aspect ratio")
|
|
image_prompt = ", ".join(prompt_parts)
|
|
|
|
# Generate image
|
|
provider = "wavespeed"
|
|
model = "ideogram-v3-turbo"
|
|
if base_avatar_bytes:
|
|
logger.info(f"[YouTube] Using character-consistent generation for scene {request.scene_id}")
|
|
style = request.style or "Realistic"
|
|
rendering_speed = request.rendering_speed or "Quality"
|
|
aspect_ratio = request.aspect_ratio or "16:9"
|
|
width = request.width or 1024
|
|
height = request.height or 576
|
|
|
|
wavespeed_client = WaveSpeedClient()
|
|
image_bytes = wavespeed_client.generate_character_image(
|
|
prompt=image_prompt,
|
|
reference_image_bytes=base_avatar_bytes,
|
|
style=style,
|
|
aspect_ratio=aspect_ratio,
|
|
rendering_speed=rendering_speed,
|
|
timeout=None,
|
|
)
|
|
model = "ideogram-character"
|
|
else:
|
|
logger.info(f"[YouTube] Generating scene {request.scene_id} from scratch")
|
|
image_options = {
|
|
"provider": "wavespeed",
|
|
"model": "ideogram-v3-turbo",
|
|
"width": request.width or 1024,
|
|
"height": request.height or 576,
|
|
}
|
|
result = generate_image(
|
|
prompt=request.custom_prompt or image_prompt,
|
|
options=image_options,
|
|
user_id=user_id,
|
|
)
|
|
image_bytes = result.image_bytes
|
|
provider = result.provider
|
|
model = result.model
|
|
|
|
# Save image
|
|
saved = _save_scene_image(image_bytes, request.scene_id)
|
|
|
|
# Save to asset library
|
|
try:
|
|
save_asset_to_library(
|
|
db=db,
|
|
user_id=user_id,
|
|
asset_type="image",
|
|
source_module="youtube_creator",
|
|
filename=saved["image_filename"],
|
|
file_url=saved["image_url"],
|
|
file_path=saved["image_path"],
|
|
file_size=len(image_bytes),
|
|
mime_type="image/png",
|
|
title=f"YouTube Scene: {request.scene_title or request.scene_id}",
|
|
description=request.scene_content or f"Scene image for {request.scene_id}",
|
|
prompt=image_prompt,
|
|
tags=["youtube_creator", "scene", request.scene_id],
|
|
provider=provider,
|
|
model=model,
|
|
asset_metadata={
|
|
"scene_id": request.scene_id,
|
|
"scene_title": request.scene_title,
|
|
"has_base_avatar": bool(base_avatar_bytes),
|
|
"width": request.width or 1024,
|
|
"height": request.height or 576,
|
|
},
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"[YouTube] Failed to save scene image to asset library: {e}")
|
|
|
|
return {
|
|
"scene_id": request.scene_id,
|
|
"scene_title": request.scene_title,
|
|
"image_filename": saved["image_filename"],
|
|
"image_url": saved["image_url"],
|
|
"width": request.width or 1024,
|
|
"height": request.height or 576,
|
|
}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as exc:
|
|
logger.error(f"[YouTube] Scene image generation failed: {exc}", exc_info=True)
|
|
raise HTTPException(status_code=500, detail=f"Failed to generate scene image: {str(exc)}")
|
|
|
|
|
|
@router.get("/images/{category}/{filename}")
|
|
async def serve_youtube_image(
|
|
category: str,
|
|
filename: str,
|
|
current_user: Dict[str, Any] = Depends(get_current_user),
|
|
):
|
|
"""
|
|
Serve stored YouTube images (avatars or scenes).
|
|
Unified endpoint for both avatar and scene images.
|
|
"""
|
|
require_authenticated_user(current_user)
|
|
|
|
if category not in {"avatars", "scenes"}:
|
|
raise HTTPException(status_code=400, detail="Invalid image category. Must be 'avatars' or 'scenes'")
|
|
|
|
if ".." in filename or "/" in filename or "\\" in filename:
|
|
raise HTTPException(status_code=400, detail="Invalid filename")
|
|
|
|
directory = YOUTUBE_AVATARS_DIR if category == "avatars" else YOUTUBE_IMAGES_DIR
|
|
image_path = directory / filename
|
|
|
|
if not image_path.exists() or not image_path.is_file():
|
|
raise HTTPException(status_code=404, detail="Image not found")
|
|
|
|
return FileResponse(
|
|
path=str(image_path),
|
|
media_type="image/png",
|
|
filename=filename,
|
|
)
|