"""YouTube Creator scene image generation handlers.""" from pathlib import Path from typing import Dict, Any, Optional import uuid from concurrent.futures import ThreadPoolExecutor from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks 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, generate_character_image from utils.asset_tracker import save_asset_to_library from utils.logger_utils import get_service_logger from ..task_manager import task_manager 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" # Thread pool for background image generation _image_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="youtube_image") 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" model: Optional[str] = None # e.g., "ideogram-v3-turbo", "qwen-image" 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) -> Optional[bytes]: """Load base avatar bytes for character consistency.""" try: # Handle different avatar URL formats if avatar_url.startswith("/api/youtube/avatars/"): # YouTube avatar filename = avatar_url.split("/")[-1].split("?")[0] avatar_path = YOUTUBE_AVATARS_DIR / filename elif avatar_url.startswith("/api/podcast/avatars/"): # Podcast avatar (cross-module usage) filename = avatar_url.split("/")[-1].split("?")[0] from pathlib import Path podcast_avatars_dir = Path(__file__).parent.parent.parent.parent / "podcast_avatars" avatar_path = podcast_avatars_dir / filename else: # Try to extract filename and check YouTube avatars first filename = avatar_url.split("/")[-1].split("?")[0] avatar_path = YOUTUBE_AVATARS_DIR / filename if not avatar_path.exists(): # Fallback to podcast avatars podcast_avatars_dir = Path(__file__).parent.parent.parent.parent / "podcast_avatars" avatar_path = podcast_avatars_dir / filename if not avatar_path.exists() or not avatar_path.is_file(): logger.warning(f"[YouTube] Avatar file not found: {avatar_path}") return None logger.info(f"[YouTube] Successfully loaded avatar: {avatar_path}") return avatar_path.read_bytes() except Exception as e: logger.error(f"[YouTube] Error loading avatar from {avatar_url}: {e}") return None 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, } class YouTubeImageTaskResponse(BaseModel): success: bool task_id: str message: str @router.post("/image", response_model=YouTubeImageTaskResponse) async def generate_youtube_scene_image( background_tasks: BackgroundTasks, request: YouTubeImageRequest, current_user: Dict[str, Any] = Depends(get_current_user), db: Session = Depends(get_db), ): """Generate a YouTube scene image with background task processing.""" logger.info(f"[YouTube] Image generation request received: scene='{request.scene_title}', user={current_user.get('id')}") user_id = require_authenticated_user(current_user) logger.info(f"[YouTube] User authenticated: {user_id}") 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}") # Create background task logger.info(f"[YouTube] Creating task for user {user_id}") task_id = task_manager.create_task("youtube_image_generation") logger.info( f"[YouTube] Created image generation task {task_id} for user {user_id}, " f"scene='{request.scene_title}'" ) # Verify task was created initial_status = task_manager.get_task_status(task_id) if not initial_status: logger.error(f"[YouTube] Failed to create task {task_id} - task not found immediately after creation") return YouTubeImageTaskResponse( success=False, task_id="", message="Failed to create image generation task. Please try again." ) # Add background task (pass request data, not database session) try: background_tasks.add_task( _execute_image_generation_task, task_id=task_id, request_data=request.dict(), # Convert to dict for background task user_id=user_id, ) logger.info(f"[YouTube] Background image generation task added for task {task_id}") except Exception as bg_error: logger.error(f"[YouTube] Failed to add background task for {task_id}: {bg_error}", exc_info=True) # Mark task as failed task_manager.update_task_status( task_id, "failed", error=str(bg_error), message="Failed to start image generation task" ) return YouTubeImageTaskResponse( success=False, task_id="", message=f"Failed to start image generation task: {str(bg_error)}" ) logger.info(f"[YouTube] Returning success response for task {task_id}") return YouTubeImageTaskResponse( success=True, task_id=task_id, message=f"Image generation started for '{request.scene_title}'" ) except HTTPException: raise except Exception as exc: logger.error(f"[YouTube] Failed to create image generation task: {exc}", exc_info=True) raise HTTPException(status_code=500, detail=f"Failed to start image generation: {str(exc)}") def _execute_image_generation_task(task_id: str, request_data: dict, user_id: str): """Background task to generate YouTube scene image.""" # Reconstruct request object from dict request = YouTubeImageRequest(**request_data) logger.info( f"[YouTubeImageGen] Background task started for task {task_id}, " f"scene='{request.scene_title}', user={user_id}" ) db = None try: # Update task status to processing task_manager.update_task_status( task_id, "processing", progress=10.0, message="Preparing image generation..." ) # Get database session for this background task from services.database import get_db db = next(get_db()) logger.info(f"[YouTubeImageGen] Database session acquired for task {task_id}") # Load avatar if provided base_avatar_bytes = None if request.base_avatar_url: base_avatar_bytes = _load_base_avatar_bytes(request.base_avatar_url) if base_avatar_bytes: logger.info(f"[YouTubeImageGen] Loaded base avatar for task {task_id}") else: logger.warning(f"[YouTubeImageGen] Could not load base avatar for task {task_id}") # Build prompt (same logic as before) 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) task_manager.update_task_status( task_id, "processing", progress=30.0, message="Generating image..." ) logger.info(f"[YouTubeImageGen] Starting image generation for task {task_id}") # Generate image (same logic as before) provider = "wavespeed" model = "ideogram-v3-turbo" if base_avatar_bytes: logger.info(f"[YouTubeImageGen] Using character-consistent generation for task {task_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 try: # Use centralized character image generation with subscription checks and tracking image_bytes = generate_character_image( prompt=image_prompt, reference_image_bytes=base_avatar_bytes, user_id=user_id, style=style, aspect_ratio=aspect_ratio, rendering_speed=rendering_speed, timeout=60, ) model = "ideogram-character" logger.info(f"[YouTubeImageGen] Character image generation successful for task {task_id}") except Exception as char_error: logger.warning(f"[YouTubeImageGen] Character generation failed for task {task_id}: {char_error}") logger.info(f"[YouTubeImageGen] Falling back to regular image generation for task {task_id}") # Fall back to regular image generation with subscription tracking image_options = { "provider": "wavespeed", "model": request.model or "ideogram-v3-turbo", "width": width, "height": height, } result = generate_image( prompt=image_prompt, options=image_options, user_id=user_id, ) image_bytes = result.image_bytes else: logger.info(f"[YouTubeImageGen] Generating scene from scratch for task {task_id}") # Use centralized image generation with subscription tracking image_options = { "provider": "wavespeed", "model": request.model or "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 # Validate image bytes before saving if not image_bytes or len(image_bytes) == 0: raise ValueError("Image generation returned empty bytes") # Basic validation: check if it's a valid image (PNG/JPEG header) if not (image_bytes.startswith(b'\x89PNG') or image_bytes.startswith(b'\xff\xd8\xff')): logger.warning(f"[YouTubeImageGen] Generated image may not be valid PNG/JPEG for task {task_id}") # Don't fail - some formats might be valid, but log warning task_manager.update_task_status( task_id, "processing", progress=80.0, message="Saving image..." ) # Save image with validation try: image_metadata = _save_scene_image(image_bytes, request.scene_id) # Verify file was saved correctly from pathlib import Path saved_path = Path(image_metadata["image_path"]) if not saved_path.exists() or saved_path.stat().st_size == 0: raise IOError(f"Image file was not saved correctly: {saved_path}") logger.info(f"[YouTubeImageGen] Image saved successfully: {saved_path} ({saved_path.stat().st_size} bytes)") except Exception as save_error: logger.error(f"[YouTubeImageGen] Failed to save image for task {task_id}: {save_error}", exc_info=True) raise # Save to asset library try: save_asset_to_library( db=db, user_id=user_id, asset_type="image", source_module="youtube_creator", filename=image_metadata["image_filename"], file_url=image_metadata["image_url"], file_path=image_metadata["image_path"], file_size=len(image_bytes), mime_type="image/png", title=f"{request.scene_title} - YouTube Scene", description=f"YouTube scene image for: {request.scene_title}", tags=["youtube_creator", "scene_image", f"scene_{request.scene_id}"], provider=provider, model=model, cost=0.10 if model == "ideogram-v3-turbo" else 0.05, asset_metadata={ "scene_id": request.scene_id, "scene_title": request.scene_title, "generation_type": "character" if base_avatar_bytes else "scene", "width": request.width or 1024, "height": request.height or 576, }, ) except Exception as e: logger.warning(f"[YouTubeImageGen] Failed to save image asset to library: {e}") # Success! task_manager.update_task_status( task_id, "completed", progress=100.0, message=f"Image generated successfully for '{request.scene_title}'", result={ "scene_id": request.scene_id, "scene_title": request.scene_title, "image_filename": image_metadata["image_filename"], "image_url": image_metadata["image_url"], "provider": provider, "model": model, "width": request.width or 1024, "height": request.height or 576, "file_size": len(image_bytes), "cost": 0.10 if model == "ideogram-v3-turbo" else 0.05, } ) logger.info(f"[YouTubeImageGen] ✅ Task {task_id} completed successfully") except Exception as exc: error_msg = str(exc) logger.error(f"[YouTubeImageGen] Task {task_id} failed: {error_msg}", exc_info=True) task_manager.update_task_status( task_id, "failed", error=error_msg, message=f"Image generation failed: {error_msg}" ) finally: if db: db.close() logger.info(f"[YouTubeImageGen] Database session closed for task {task_id}") @router.get("/image/status/{task_id}") async def get_image_generation_status( task_id: str, current_user: Dict[str, Any] = Depends(get_current_user), ): """ Get the status of an image generation task. Returns current progress, status, and result when complete. """ require_authenticated_user(current_user) logger.info(f"[YouTubeAPI] Getting image generation status for task: {task_id}") task_status = task_manager.get_task_status(task_id) if task_status: logger.info(f"[YouTubeAPI] Task {task_id} status: {task_status.get('status', 'unknown')}, progress: {task_status.get('progress', 0)}, has_result: {'result' in task_status}") if not task_status: logger.warning( f"[YouTubeAPI] Image generation task {task_id} not found." ) raise HTTPException( status_code=404, detail={ "error": "Task not found", "message": "The image generation task was not found. It may have expired, been cleaned up, or the server may have restarted.", "task_id": task_id, "user_action": "Please try generating the image again." } ) return task_status @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, )