Base code
This commit is contained in:
Kunthawat Greethong
470
backend/api/youtube/handlers/images.py
Normal file
470
backend/api/youtube/handlers/images.py
Normal file
@@ -0,0 +1,470 @@
|
||||
"""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,
|
||||
)
|
||||
Reference in New Issue
Block a user