refactor: update agent and server code, move agent.py to root

agent.py

@@ -1,21 +1,23 @@
 import os
+import time
+import shutil
+import requests
+import tempfile
+import logging
 from typing import TypedDict, Optional
 from langgraph.graph import StateGraph, END
-from langchain_google_genai import ChatGoogleGenerativeAI
 from google import genai
+from groq import Groq
 from gradio_client import Client, handle_file
-import shutil
-import requests
-import tempfile
-import os
-import shutil
-import requests
-import tempfile
-
 from dotenv import load_dotenv
 
+# Load environment variables
 load_dotenv()
 
+# Configure Logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # State Definition
 class ContinuityState(TypedDict):
     video_a_url: str
@@ -27,27 +29,25 @@ class ContinuityState(TypedDict):
     video_a_local_path: Optional[str]
     video_c_local_path: Optional[str]
 
-# Node 1: Analyst
+# --- HELPER FUNCTIONS ---
+def download_to_temp(url):
+    logger.info(f"Downloading: {url}")
+    resp = requests.get(url, stream=True)
+    resp.raise_for_status()
+    suffix = os.path.splitext(url.split("/")[-1])[1] or ".mp4"
+    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as f:
+        shutil.copyfileobj(resp.raw, f)
+        return f.name
+
+# --- NODE 1: ANALYST ---
 def analyze_videos(state: ContinuityState) -> dict:
-    print("--- 🧐 Analyst Node (Director) ---")
+    logger.info("--- 🧐 Analyst Node (Director) ---")
     
     video_a_url = state['video_a_url']
     video_c_url = state['video_c_url']
     
-    # Initialize Google GenAI Client
-    client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
-    
+    # 1. Prepare Files
     try:
-        # Download videos to temp files for analysis
-        def download_to_temp(url):
-            print(f"Downloading: {url}")
-            resp = requests.get(url, stream=True)
-            resp.raise_for_status()
-            suffix = os.path.splitext(url.split("/")[-1])[1] or ".mp4"
-            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as f:
-                shutil.copyfileobj(resp.raw, f)
-                return f.name
-
         path_a = state.get('video_a_local_path')
         if not path_a:
              path_a = download_to_temp(video_a_url)
@@ -55,94 +55,103 @@ def analyze_videos(state: ContinuityState) -> dict:
         path_c = state.get('video_c_local_path')
         if not path_c:
              path_c = download_to_temp(video_c_url)
-        
-        print("Uploading videos to Gemini...")
-        file_a = client.files.upload(file=path_a)
-        file_c = client.files.upload(file=path_c)
-        
-        # Wait for processing? Usually quick for small files, but good practice to check state if needed.
-        # For simplicity in this agent, assuming ready or waiting implicitly. 
-        # (Gemini 1.5 Flash usually processes quickly)
-
-        prompt = """
-        You are a film director. 
-        Analyze the motion, lighting, and subject of the first video (Video A) and the second video (Video C). 
-        Write a detailed visual prompt for a 2-second video (Video B) that smoothly transitions from the end of A to the start of C.
-        Target Output: A single concise descriptive paragraph for the video generation model.
-        """
-        
-        print("Generating transition prompt...")
-        response = client.models.generate_content(
-            model="gemini-1.5-flash",
-            contents=[prompt, file_a, file_c]
-        )
-        
-        transition_prompt = response.text
-        print(f"Generated Prompt: {transition_prompt}")
-        
-        # Cleanup uploaded files from local ? (Files on server stay for 48h or until deleted)
-        # client.files.delete(name=file_a.name) 
-        # client.files.delete(name=file_c.name)
-        
-        # We also need these local paths for the Generator node to extract frames!
-        # Pass them in state or re-download? Better to pass paths if possible, but 
-        # State definition expects URLs. We can add temp paths to state or re-download.
-        # Let's add temp paths to state for efficiency.
-        
-        return {
-            "scene_analysis": transition_prompt, 
-            "veo_prompt": transition_prompt,
-            "video_a_local_path": path_a,
-            "video_c_local_path": path_c
-        }
-
     except Exception as e:
-        print(f"Error in Analyst: {e}")
-        return {"scene_analysis": f"Error: {str(e)}", "veo_prompt": "Error"}
+        logger.error(f"Download failed: {e}")
+        return {"scene_analysis": "Error downloading", "veo_prompt": "Smooth cinematic transition"}
 
+    # 2. Try Gemini 2.0 (With Retry)
+    client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+    transition_prompt = None
+    
+    retries = 3
+    for attempt in range(retries):
+        try:
+            logger.info(f"Uploading videos to Gemini... (Attempt {attempt+1})")
+            file_a = client.files.upload(file=path_a)
+            file_c = client.files.upload(file=path_c)
+            
+            prompt_text = """
+            You are a film director. 
+            Analyze the motion, lighting, and subject of the first video (Video A) and the second video (Video C). 
+            Write a detailed visual prompt for a 2-second video (Video B) that smoothly transitions from the end of A to the start of C.
+            Target Output: A single concise descriptive paragraph for the video generation model.
+            """
+            
+            logger.info("Generating transition prompt...")
+            # Using 2.0 Flash as per your logs (or 1.5-flash if preferred)
+            response = client.models.generate_content(
+                model="gemini-2.0-flash-exp", 
+                contents=[prompt_text, file_a, file_c]
+            )
+            transition_prompt = response.text
+            logger.info(f"Generated Prompt: {transition_prompt}")
+            break # Success
+
+        except Exception as e:
+            if "429" in str(e) or "RESOURCE_EXHAUSTED" in str(e):
+                wait = 30 * (attempt + 1)
+                logger.warning(f"⚠️ Gemini Quota 429. Retrying in {wait}s...")
+                time.sleep(wait)
+            else:
+                logger.error(f"⚠️ Gemini Error: {e}")
+                break
 
-# Node 2: Generator (Wan 2.2 First Last Frame)
+    # 3. Fallback: Groq (If Gemini failed)
+    if not transition_prompt:
+        logger.info("Switching to Llama 3.2 (Groq) Fallback...")
+        try:
+            groq_client = Groq(api_key=os.environ["GROQ_API_KEY"])
+            # We can't easily send videos, so we generate a prompt based on general best practices
+            fallback_prompt = "Create a smooth, cinematic visual transition that bridges two scenes with matching lighting and motion blur."
+            
+            completion = groq_client.chat.completions.create(
+                model="llama-3.2-11b-vision-preview",
+                messages=[
+                    {"role": "user", "content": f"Refine this into a video generation prompt: {fallback_prompt}"}
+                ]
+            )
+            transition_prompt = completion.choices[0].message.content
+        except Exception as e:
+            logger.error(f"❌ Groq also failed: {e}")
+            transition_prompt = "Smooth cinematic transition with motion blur matching the scenes."
+
+    return {
+        "scene_analysis": transition_prompt, 
+        "veo_prompt": transition_prompt,
+        "video_a_local_path": path_a,
+        "video_c_local_path": path_c
+    }
+
+# --- NODE 2: GENERATOR ---
 def generate_video(state: ContinuityState) -> dict:
-    print("--- 🎥 Generator Node (Wan 2.2) ---")
+    logger.info("--- 🎥 Generator Node ---")
     
     prompt = state.get('veo_prompt', "")
     path_a = state.get('video_a_local_path')
     path_c = state.get('video_c_local_path')
     
     if not path_a or not path_c:
-        # Fallback if dependencies failed or state clean
-        # Re-download logic would go here, but assuming flow works
         return {"generated_video_url": "Error: Missing local video paths"}
 
     try:
-        # Extract Frames
+        # Extract Frames (simplified for brevity, ensuring libraries are imported)
         import cv2
         from PIL import Image
         
         def get_frame(video_path, location="last"):
             cap = cv2.VideoCapture(video_path)
             total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            if location == "last":
-                cap.set(cv2.CAP_PROP_POS_FRAMES, total_frames - 1)
-            else: # first
-                cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
-            
+            if location == "last": cap.set(cv2.CAP_PROP_POS_FRAMES, total_frames - 1)
+            else: cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
             ret, frame = cap.read()
             cap.release()
-            
-            if ret:
-                # Convert BGR to RGB
-                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                return Image.fromarray(frame_rgb)
-            else:
-                raise ValueError(f"Could not extract frame from {video_path}")
+            if ret: return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            raise ValueError(f"Could not extract frame from {video_path}")
 
-        print("Extracting frames...")
+        logger.info("Extracting frames...")
         img_start = get_frame(path_a, "last")
         img_end = get_frame(path_c, "first")
         
-        # Save frames to temp files for Gradio Client (it handles file paths better than PIL objects usually)
-        # Although client.predict might take PIL, handle_file is safer with paths.
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as f_start:
             img_start.save(f_start, format="PNG")
             start_path = f_start.name
@@ -151,61 +160,64 @@ def generate_video(state: ContinuityState) -> dict:
             img_end.save(f_end, format="PNG")
             end_path = f_end.name
 
-        # Call Wan 2.2
-        print("Initializing Wan Client...")
-        client = Client("multimodalart/wan-2-2-first-last-frame")
-        
-        print(f"Generating transition with prompt: {prompt[:50]}...")
-        # predict(start_image, end_image, prompt, negative_prompt, duration, steps, guide, guide2, seed, rand, api_name)
-        result = client.predict(
-            start_image_pil=handle_file(start_path),
-            end_image_pil=handle_file(end_path),
-            prompt=prompt,
-            negative_prompt="blurry, distorted, low quality, static",
-            duration_seconds=2.1,
-            steps=20, # Default is often around 20-30 for good quality
-            guidance_scale=5.0,
-            guidance_scale_2=5.0,
-            seed=42,
-            randomize_seed=True,
-            api_name="/generate_video"
-        )
-        
-        # Clean up temp frames and videos
+        # --- ATTEMPT 1: WAN 2.2 ---
         try:
-            os.remove(start_path)
-            os.remove(end_path)
-            os.remove(path_a)
-            os.remove(path_c)
-        except:
-            pass
-
-        # Parse valid result
-        # Expected: ({'video': path, ...}, seed) or just path depending on version
-        # Based on inspection: (generated_video_mp4, seed)
-        video_out = result[0]
-        if isinstance(video_out, dict) and 'video' in video_out:
-             return {"generated_video_url": video_out['video']}
-        elif isinstance(video_out, str) and os.path.exists(video_out):
-             return {"generated_video_url": video_out}
-        else:
-             return {"generated_video_url": f"Error: Unexpected output {result}"}
+            logger.info("Initializing Wan Client...")
+            client = Client("multimodalart/wan-2-2-first-last-frame")
+            
+            logger.info(f"Generating with Wan 2.2... Prompt: {prompt[:30]}...")
+            result = client.predict(
+                start_image_pil=handle_file(start_path),
+                end_image_pil=handle_file(end_path),
+                prompt=prompt,
+                negative_prompt="blurry, distorted, low quality, static",
+                duration_seconds=2.1,
+                steps=20,
+                guidance_scale=5.0,
+                guidance_scale_2=5.0,
+                seed=42,
+                randomize_seed=True,
+                api_name="/generate_video"
+            )
+            # Handle Wan output format
+            video_out = result[0]
+            if isinstance(video_out, dict) and 'video' in video_out:
+                 return {"generated_video_url": video_out['video']}
+            elif isinstance(video_out, str) and os.path.exists(video_out):
+                 return {"generated_video_url": video_out}
+                 
+        except Exception as e:
+            logger.warning(f"⚠️ Wan 2.2 Failed: {e}")
+
+        # --- ATTEMPT 2: SVD FALLBACK ---
+        logger.info("🔄 Switching to SVD Fallback...")
+        try:
+            # FIXED REPO ID
+            client = Client("multimodalart/stable-video-diffusion")
+            
+            # SVD uses one image, we'll use the start frame
+            result = client.predict(
+                handle_file(start_path),
+                0.0, 0.0, 1, 25, # resized_width, resized_height, motion_bucket_id, fps
+                api_name="/predict"
+            )
+            logger.info(f"✅ SVD Generated: {result}")
+            return {"generated_video_url": result} # SVD usually returns path string
+            
+        except Exception as e:
+            logger.error(f"❌ All Generators Failed. Error: {e}")
+            return {"generated_video_url": f"Error: {str(e)}"}
 
     except Exception as e:
-        print(f"Error in Generator: {e}")
+        logger.error(f"Error in Generator Setup: {e}")
         return {"generated_video_url": f"Error: {str(e)}"}
 
 
 # Graph Construction
 workflow = StateGraph(ContinuityState)
-
 workflow.add_node("analyst", analyze_videos)
-# workflow.add_node("prompter", draft_prompt) # Skipped, Analyst does extraction + prompting
 workflow.add_node("generator", generate_video)
-
 workflow.set_entry_point("analyst")
-
 workflow.add_edge("analyst", "generator")
 workflow.add_edge("generator", END)
-
-app = workflow.compile()
+app = workflow.compile()

continuity_agent/agent.py

@@ -1,256 +0,0 @@
-import os
-import time
-import shutil
-import cv2
-import numpy as np
-import base64
-import tempfile
-from groq import Groq
-from google import genai
-from gradio_client import Client, handle_file
-from dotenv import load_dotenv
-
-load_dotenv()
-
-# --- HELPER: Filmstrip Engine ---
-def create_filmstrip(video_path, samples=5, is_start=False):
-    """Extracts frames and stitches them into a filmstrip for Vision analysis."""
-    try:
-        cap = cv2.VideoCapture(video_path)
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        duration = total_frames / fps
-
-        # Determine extraction points
-        if is_start: # First 2 seconds
-            start_f = 0
-            end_f = int(min(total_frames, 2 * fps))
-            if end_f <= start_f: end_f = total_frames # Handle short videos
-        else: # Last 2 seconds
-            start_f = int(max(0, total_frames - 2 * fps))
-            end_f = total_frames
-            if start_f >= end_f: start_f = 0
-            
-        frame_indices = np.linspace(start_f, end_f - 1, samples, dtype=int)
-        frames = []
-        
-        for idx in frame_indices:
-            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
-            ret, frame = cap.read()
-            if ret:
-                # Resize for token efficiency (Height 300px)
-                h, w, _ = frame.shape
-                scale = 300 / h
-                new_w = int(w * scale)
-                frame = cv2.resize(frame, (new_w, 300))
-                frames.append(frame)
-        cap.release()
-        
-        if not frames:
-            raise ValueError("No frames extracted")
-            
-        # Stitch horizontally
-        filmstrip = cv2.hconcat(frames)
-        
-        # Use a consistent temp file pattern or unique name
-        temp_dir = tempfile.gettempdir()
-        output_path = os.path.join(temp_dir, f"temp_strip_{int(time.time())}_{'start' if is_start else 'end'}.jpg")
-        cv2.imwrite(output_path, filmstrip)
-        return output_path
-    except Exception as e:
-        print(f"⚠️ Filmstrip failed: {e}")
-        return None
-
-# --- PHASE 1: ANALYZE ONLY ---
-def analyze_only(video_a_path: str, video_c_path: str):
-    print(f"🎬 Analyst: Processing videos...")
-
-    # Generate Filmstrips
-    strip_a = create_filmstrip(video_a_path, is_start=False)
-    strip_c = create_filmstrip(video_c_path, is_start=True)
-    
-    if not strip_a or not strip_c:
-        return {
-            "prompt": "Cinematic transition between scenes.",
-            "video_a_path": video_a_path,
-            "video_c_path": video_c_path,
-            "status": "warning",
-            "detail": "Could not create filmstrips"
-        }
-
-    prompt = "Smooth cinematic transition." # Default safety
-
-    # 1. Try Gemini 2.0 (Primary)
-    try:
-        print("🤖 Engaging Gemini 2.0...")
-        client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY")) # Using correct env var name
-        
-        file_a = client.files.upload(file=strip_a)
-        file_c = client.files.upload(file=strip_c)
-
-        system_prompt = """
-        You are an expert film editor. Analyze these two 'filmstrips'. 
-        Image 1 shows the end of the first clip (time flows left-to-right). 
-        Image 2 shows the start of the next clip. 
-        Describe the motion, lighting, and subject connection required to seamlessly bridge A to C in a cinematic way.
-        Output a SINGLE concise paragraph for a video generation model.
-        """
-        
-        response = client.models.generate_content(
-            model="gemini-2.0-flash",
-            contents=[system_prompt, file_a, file_c]
-        )
-        
-        if response.text:
-            prompt = response.text
-            
-        # raise Exception("Force Fallback for Testing") # Commented out for production use unless specifically testing
-        
-    except Exception as e:
-        print(f"⚠️ Gemini Quota/Error: {e}. Switching to Llama 3.2 (Groq)...")
-        
-        # 2. Try Groq (Fallback)
-        try:
-            groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-            
-            def encode_image(image_path):
-                with open(image_path, "rb") as image_file:
-                    return base64.b64encode(image_file.read()).decode('utf-8')
-            
-            b64_a = encode_image(strip_a)
-            b64_c = encode_image(strip_c)
-            
-            completion = groq_client.chat.completions.create(
-                model="llama-3.2-11b-vision-instruct", 
-                messages=[
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "text", "text": "These images show the END of Clip A and START of Clip C. Describe a smooth visual transition to bridge them."},
-                            {
-                                "type": "image_url", 
-                                "image_url": {"url": f"data:image/jpeg;base64,{b64_a}"}
-                            },
-                             {
-                                "type": "image_url", 
-                                "image_url": {"url": f"data:image/jpeg;base64,{b64_c}"}
-                            }
-                        ]
-                    }
-                ],
-                temperature=0.7,
-                max_tokens=500
-            )
-            prompt = completion.choices[0].message.content
-        except Exception as groq_e:
-            print(f"❌ Groq also failed: {groq_e}. Using default prompt.")
-
-    # Cleanup
-    try:
-        if os.path.exists(strip_a): os.remove(strip_a)
-        if os.path.exists(strip_c): os.remove(strip_c)
-    except:
-        pass
-
-    return {
-        "prompt": prompt,
-        "video_a_path": video_a_path,
-        "video_c_path": video_c_path,
-        "status": "success"
-    }
-
-# --- PHASE 2: GENERATE ONLY ---
-def generate_only(prompt: str, video_a_path: str, video_c_path: str):
-    print(f"🎥 Generator: Action! Prompt: {prompt[:50]}...")
-
-    # 1. Primary: Wan 2.2
-    try:
-        # Extract Frames for Wan
-        # We need to save temporary frames because handle_file expects a path
-        def get_frame(v_path, at_start):
-            cap = cv2.VideoCapture(v_path)
-            if not at_start:
-                 total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-                 cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, total-1))
-            ret, frame = cap.read()
-            cap.release()
-            if not ret: raise ValueError("Frame extract failed")
-            
-            # Resize safe for Wan
-            h, w = frame.shape[:2]
-            if h > 480:
-                scale = 480/h
-                frame = cv2.resize(frame, (int(w*scale), 480))
-            
-            t_path = os.path.join(tempfile.gettempdir(), f"wan_frame_{int(time.time())}_{'s' if at_start else 'e'}.png")
-            cv2.imwrite(t_path, frame)
-            return t_path
-
-        f_start = get_frame(video_a_path, False) # Last frame of A
-        f_end = get_frame(video_c_path, True)    # First frame of C
-
-        client = Client("multimodalart/wan-2-2-first-last-frame", token=os.environ.get("HF_TOKEN"))
-        
-        print("Generating with Wan 2.2...")
-        result = client.predict(
-            start_image_pil=handle_file(f_start),
-            end_image_pil=handle_file(f_end),
-            prompt=prompt,
-            negative_prompt="blurry, distorted, low quality, static",
-            duration_seconds=2.1,
-            steps=20,
-            guidance_scale=5.0,
-            guidance_scale_2=5.0,
-            seed=42,
-            randomize_seed=True,
-            api_name="/generate_video" 
-        )
-        
-        # Cleanup temp
-        try:
-            os.remove(f_start)
-            os.remove(f_end)
-        except: pass
-
-        # Parse result
-        video_out = result[0]
-        if isinstance(video_out, dict) and 'video' in video_out:
-             return {"video_url": video_out['video']}
-        elif isinstance(video_out, str):
-             return {"video_url": video_out}
-        else:
-             raise ValueError(f"Unknown Wan output: {result}")
-
-    except Exception as e:
-        print(f"⚠️ Wan 2.2 Failed (Quota/Error): {e}")
-        print("🔄 Switching to SVD Fallback...")
-        
-        # 2. Fallback: SVD (Image-to-Video)
-        try:
-            client_svd = Client("stabilityai/stable-video-diffusion-img2vid-xt-1-1")
-            
-            # Extract last frame of A for SVD input
-            cap = cv2.VideoCapture(video_a_path)
-            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            cap.set(cv2.CAP_PROP_POS_FRAMES, total-1)
-            ret, frame = cap.read()
-            cap.release()
-            
-            # Resize for SVD (1024x576 recommended or similar 16:9)
-            frame = cv2.resize(frame, (1024, 576))
-            
-            svd_input_path = os.path.join(tempfile.gettempdir(), "svd_input.jpg")
-            cv2.imwrite(svd_input_path, frame)
-            
-            print("Generating with SVD...")
-            result = client_svd.predict(
-                svd_input_path,
-                0.0, 127, 6, 
-                api_name="/predict"
-            )
-            return {"video_url": result} 
-            
-        except Exception as svd_e:
-            err_msg = f"All Generators Failed. Wan: {e}, SVD: {svd_e}"
-            print(f"❌ {err_msg}")
-            return {"video_url": f"Error: {err_msg}"}

server.py

@@ -71,7 +71,7 @@ async def generate_endpoint(
              
         # Call Agent
         result = generate_only(prompt, video_a_path, video_c_path)
-        gen_path = result.get("video_url")
+        gen_path = result.get("generated_video_url")
         
         if not gen_path or "Error" in gen_path:
             raise HTTPException(status_code=500, detail=f"Generation failed: {gen_path}")