"""Video generation service using Replicate API with R2 storage.""" import asyncio import json import os import time import uuid import httpx from loguru import logger from pipecat.adapters.schemas.function_schema import FunctionSchema from pipecat.frames.frames import OutputTransportMessageUrgentFrame from pipecat.services.llm_service import FunctionCallParams from app.services.r2 import upload_bytes from typing import TYPE_CHECKING if TYPE_CHECKING: from bot.core.config import BotSessionState, TranscriptManager def _get_replicate_token(): return os.getenv("REPLICATE_API_TOKEN", "") REPLICATE_VIDEO_MODEL = "wan-video/wan-2.2-t2v-fast" REPLICATE_I2V_MODEL = "wan-video/wan-2.2-i2v-fast" # --- Function Schemas --- generate_video_function = FunctionSchema( name="generate_video", description=( "Generate a video from a text description using AI. Use when user asks to create a VIDEO " "WITHOUT having uploaded an image. Do NOT call if video generation is already in progress." ), properties={ "prompt": {"type": "string", "description": "Detailed description of the video to generate"}, }, required=["prompt"], ) generate_video_from_image_function = FunctionSchema( name="generate_video_from_image", description=( "Generate a video FROM an uploaded/generated image. Use when user has an image AND wants a video from it. " "Do NOT call if video generation is already in progress." ), properties={ "prompt": {"type": "string", "description": "Description of how to animate the image"}, }, required=["prompt"], ) # --- Helpers --- async def _send_rtvi(task, data: dict): try: frame = OutputTransportMessageUrgentFrame(message={ "label": "rtvi-ai", "type": "server-message", "data": data, }) await task.queue_frame(frame) except Exception as e: logger.warning(f"Failed to send RTVI message: {e}") async def _download_video(url: str) -> bytes: async with httpx.AsyncClient(timeout=120.0) as client: response = await client.get(url) response.raise_for_status() return response.content async def _run_replicate(model: str, replicate_input: dict) -> str: """Run Replicate model and return output URL.""" import replicate os.environ["REPLICATE_API_TOKEN"] = _get_replicate_token() def _run(): output = replicate.run(model, input=replicate_input) # Handle different output types: FileOutput, list, iterator, string if hasattr(output, "url"): return output.url if isinstance(output, list): item = output[0] if output else output return item.url if hasattr(item, "url") else str(item) if isinstance(output, str): return output # Iterator — consume first item try: item = next(iter(output)) return item.url if hasattr(item, "url") else str(item) except (StopIteration, TypeError): return str(output) return await asyncio.to_thread(_run) # --- Text-to-Video Handler --- def create_video_handler( bot_session: "BotSessionState", transcript_manager: "TranscriptManager", task, ): """Create text-to-video handler.""" async def handle_generate_video(params: FunctionCallParams): prompt = params.arguments.get("prompt", "") session_id = bot_session.session_id or "unknown" if not prompt: await params.result_callback(json.dumps({ "status": "error", "assistant_instruction": "No prompt provided. Ask the user what video they want.", })) return if bot_session.is_processing: await params.result_callback(json.dumps({ "status": "busy", "assistant_instruction": "A generation is already in progress. Wait silently.", })) return # Usage limit check limit_reason = bot_session.can_generate_video() if limit_reason == "subscribe": await _send_rtvi(task, { "type": "limit_reached", "data": {"media_type": "video", "reason": "subscribe"}, }) await params.result_callback(json.dumps({ "status": "limit_reached", "assistant_instruction": "The user has used all their free video credits. Tell them to subscribe to Maya Pro for more images and videos.", })) return elif limit_reason == "topup": await _send_rtvi(task, { "type": "limit_reached", "data": {"media_type": "video", "reason": "topup"}, }) await params.result_callback(json.dumps({ "status": "limit_reached", "assistant_instruction": "The user has used all their video credits for this subscription period. Tell them they can buy a top-up pack from the app to get more credits.", })) return bot_session.is_processing = True video_id = f"vid_{int(time.time())}_{uuid.uuid4().hex[:4]}" await _send_rtvi(task, {"type": "generation_started", "data": {"type": "video", "status": "in_progress"}}) try: logger.info(f"VIDEO_GENERATE_START: session={session_id} prompt={prompt[:60]}") start_time = time.time() video_url = await _run_replicate(REPLICATE_VIDEO_MODEL, {"prompt": prompt, "resolution": "480p"}) video_bytes = await _download_video(video_url) r2_key = f"sessions/{session_id}/videos/{video_id}.mp4" r2_url = await asyncio.to_thread(upload_bytes, video_bytes, r2_key, "video/mp4") total_time = time.time() - start_time logger.info(f"VIDEO_GENERATED: session={session_id} video_id={video_id} url={r2_url} time={total_time:.1f}s") bot_session.add_generated_video( video_type="text_to_video", metadata={"video_id": video_id, "r2_url": r2_url, "model": REPLICATE_VIDEO_MODEL, "total_time": total_time}, prompt=prompt, ) bot_session.use_video() await _send_rtvi(task, { "type": "video_generation_complete", "data": {"video_url": r2_url, "video_id": video_id, "description": prompt}, }) await params.result_callback(json.dumps({ "status": "completed", "video_url": r2_url, "video_id": video_id, "assistant_instruction": "Video is ready! Tell the user their video is generated and they can see it on screen.", })) except Exception as e: logger.error(f"VIDEO_ERROR: session={session_id} error={e}", exc_info=True) await _send_rtvi(task, { "type": "generation_failed", "data": {"error": str(e)[:100], "media_type": "video"}, }) await params.result_callback(json.dumps({ "status": "error", "assistant_instruction": f"Video generation failed: {str(e)[:100]}. Apologize and suggest trying again.", })) finally: bot_session.is_processing = False return handle_generate_video # --- Image-to-Video Handler --- def create_video_from_image_handler( bot_session: "BotSessionState", transcript_manager: "TranscriptManager", task, ): """Create image-to-video handler.""" async def handle_generate_video_from_image(params: FunctionCallParams): prompt = params.arguments.get("prompt", "") session_id = bot_session.session_id or "unknown" if bot_session.is_processing: await params.result_callback(json.dumps({ "status": "busy", "assistant_instruction": "A generation is already in progress. Wait silently.", })) return # Usage limit check limit_reason = bot_session.can_generate_video() if limit_reason == "subscribe": await _send_rtvi(task, { "type": "limit_reached", "data": {"media_type": "video", "reason": "subscribe"}, }) await params.result_callback(json.dumps({ "status": "limit_reached", "assistant_instruction": "The user has used all their free video credits. Tell them to subscribe to Maya Pro for more images and videos.", })) return elif limit_reason == "topup": await _send_rtvi(task, { "type": "limit_reached", "data": {"media_type": "video", "reason": "topup"}, }) await params.result_callback(json.dumps({ "status": "limit_reached", "assistant_instruction": "The user has used all their video credits for this subscription period. Tell them they can buy a top-up pack from the app to get more credits.", })) return source_image = bot_session.get_uploaded_image() if not source_image: await params.result_callback(json.dumps({ "status": "error", "assistant_instruction": "No image available. Ask the user to upload or generate an image first.", })) return image_url = source_image.get("r2_url") or source_image.get("url") if not image_url: await params.result_callback(json.dumps({ "status": "error", "assistant_instruction": "Image URL not available. Ask user to upload again.", })) return if not prompt: prompt = f"Animate this: {source_image.get('description', 'the image')}" bot_session.is_processing = True video_id = f"i2v_{int(time.time())}_{uuid.uuid4().hex[:4]}" await _send_rtvi(task, {"type": "generation_started", "data": {"type": "video", "status": "in_progress"}}) try: logger.info(f"I2V_GENERATE_START: session={session_id} prompt={prompt[:60]}") start_time = time.time() video_url = await _run_replicate(REPLICATE_I2V_MODEL, { "image": image_url, "prompt": prompt, "go_fast": True, "num_frames": 81, "resolution": "480p", "sample_shift": 12, "frames_per_second": 16, }) video_bytes = await _download_video(video_url) r2_key = f"sessions/{session_id}/videos/{video_id}.mp4" r2_url = await asyncio.to_thread(upload_bytes, video_bytes, r2_key, "video/mp4") total_time = time.time() - start_time logger.info(f"I2V_GENERATED: session={session_id} video_id={video_id} url={r2_url} time={total_time:.1f}s") bot_session.add_generated_video( video_type="image_to_video", metadata={"video_id": video_id, "r2_url": r2_url, "model": REPLICATE_I2V_MODEL, "source_image": image_url, "total_time": total_time}, prompt=prompt, ) bot_session.use_video() await _send_rtvi(task, { "type": "video_generation_complete", "data": {"video_url": r2_url, "video_id": video_id, "description": prompt}, }) await params.result_callback(json.dumps({ "status": "completed", "video_url": r2_url, "video_id": video_id, "assistant_instruction": "Video from image is ready! Tell the user they can see it on screen.", })) except Exception as e: logger.error(f"I2V_ERROR: session={session_id} error={e}", exc_info=True) await _send_rtvi(task, { "type": "generation_failed", "data": {"error": str(e)[:100], "media_type": "video"}, }) await params.result_callback(json.dumps({ "status": "error", "assistant_instruction": f"Video from image failed: {str(e)[:100]}. Apologize and suggest trying again.", })) finally: bot_session.is_processing = False return handle_generate_video_from_image