Upload agent.py

agent.py

@@ -0,0 +1,489 @@
+"""
+ViralCut Agent - Runtime
+========================
+The actual agent that uses the fine-tuned model to edit videos autonomously.
+
+This connects the trained model to real tools:
+- FFmpeg for video editing
+- DuckDuckGo for web search (free, no API key)
+- Whisper for transcription
+- PySceneDetect for shot detection
+
+Usage:
+    python agent.py --video raw_footage.mp4 --platform tiktok --niche food
+    python agent.py --plan --niche "coffee shop" --platform tiktok
+"""
+
+import argparse
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+# ============================================================
+# TOOL IMPLEMENTATIONS
+# ============================================================
+
+class FFmpegTool:
+    """Execute FFmpeg commands for video/audio processing."""
+    
+    @staticmethod
+    def run(command: str, description: str = "") -> str:
+        """Execute an FFmpeg command and return result."""
+        print(f"  🎬 FFmpeg: {description}")
+        print(f"     $ {command}")
+        try:
+            result = subprocess.run(
+                command, shell=True, capture_output=True, text=True, timeout=120
+            )
+            if result.returncode == 0:
+                return json.dumps({"status": "success", "message": f"Command completed: {description}"})
+            else:
+                return json.dumps({"status": "error", "message": result.stderr[:500]})
+        except subprocess.TimeoutExpired:
+            return json.dumps({"status": "error", "message": "Command timed out after 120s"})
+        except Exception as e:
+            return json.dumps({"status": "error", "message": str(e)})
+
+
+class WebSearchTool:
+    """Search the web using DuckDuckGo (free, no API key needed)."""
+    
+    @staticmethod
+    def search(query: str, search_type: str = "general") -> str:
+        """Search the web and return results."""
+        print(f"  🔍 Searching: {query} (type: {search_type})")
+        try:
+            from duckduckgo_search import DDGS
+            with DDGS() as ddgs:
+                results = []
+                for r in ddgs.text(query, max_results=5):
+                    results.append({
+                        "title": r.get("title", ""),
+                        "url": r.get("href", ""),
+                        "description": r.get("body", "")[:200]
+                    })
+                return json.dumps({"results": results})
+        except ImportError:
+            return json.dumps({"results": [{"title": "Install duckduckgo-search", "description": "pip install duckduckgo-search"}]})
+        except Exception as e:
+            return json.dumps({"results": [], "error": str(e)})
+
+
+class VideoAnalyzer:
+    """Analyze video files using ffprobe and PySceneDetect."""
+    
+    @staticmethod
+    def analyze(video_path: str, analysis_type: str = "full") -> str:
+        """Analyze a video file."""
+        print(f"  📊 Analyzing: {video_path} ({analysis_type})")
+        
+        if not os.path.exists(video_path):
+            return json.dumps({"error": f"File not found: {video_path}"})
+        
+        result = {}
+        
+        # Get basic info via ffprobe
+        try:
+            probe = subprocess.run(
+                f'ffprobe -v quiet -print_format json -show_format -show_streams "{video_path}"',
+                shell=True, capture_output=True, text=True
+            )
+            if probe.returncode == 0:
+                info = json.loads(probe.stdout)
+                fmt = info.get("format", {})
+                result["duration"] = float(fmt.get("duration", 0))
+                result["size_mb"] = round(int(fmt.get("size", 0)) / 1024 / 1024, 1)
+                
+                for stream in info.get("streams", []):
+                    if stream.get("codec_type") == "video":
+                        result["resolution"] = f"{stream.get('width')}x{stream.get('height')}"
+                        result["fps"] = eval(stream.get("r_frame_rate", "30/1"))
+                        result["codec"] = stream.get("codec_name")
+                    elif stream.get("codec_type") == "audio":
+                        result["audio_codec"] = stream.get("codec_name")
+                        result["audio_channels"] = stream.get("channels")
+        except Exception as e:
+            result["probe_error"] = str(e)
+        
+        # Scene detection
+        if analysis_type in ("full", "scenes"):
+            try:
+                from scenedetect import open_video, SceneManager
+                from scenedetect.detectors import ContentDetector
+                
+                video = open_video(video_path)
+                scene_manager = SceneManager()
+                scene_manager.add_detector(ContentDetector(threshold=27))
+                scene_manager.detect_scenes(video)
+                scene_list = scene_manager.get_scene_list()
+                
+                result["scenes"] = []
+                for i, (start, end) in enumerate(scene_list):
+                    result["scenes"].append({
+                        "scene": i + 1,
+                        "start": round(start.get_seconds(), 2),
+                        "end": round(end.get_seconds(), 2),
+                        "duration": round((end - start).get_seconds(), 2)
+                    })
+            except ImportError:
+                result["scenes_note"] = "Install scenedetect: pip install scenedetect[opencv]"
+            except Exception as e:
+                result["scenes_error"] = str(e)
+        
+        # Transcript via Whisper
+        if analysis_type in ("full", "transcript", "audio"):
+            try:
+                import whisper
+                model = whisper.load_model("base")
+                transcript = model.transcribe(video_path)
+                result["transcript"] = transcript.get("text", "")[:2000]
+                result["segments"] = [
+                    {"start": s["start"], "end": s["end"], "text": s["text"]}
+                    for s in transcript.get("segments", [])[:50]
+                ]
+            except ImportError:
+                result["transcript_note"] = "Install whisper: pip install openai-whisper"
+            except Exception as e:
+                result["transcript_error"] = str(e)
+        
+        return json.dumps(result)
+
+
+class ViralityScorer:
+    """Score video content for viral potential."""
+    
+    @staticmethod
+    def score(video_path: str, platform: str, niche: str = "") -> str:
+        """Score a video's viral potential based on heuristics."""
+        print(f"  📈 Scoring virality: {video_path} for {platform}")
+        
+        # Get video info
+        try:
+            probe = subprocess.run(
+                f'ffprobe -v quiet -print_format json -show_format -show_streams "{video_path}"',
+                shell=True, capture_output=True, text=True
+            )
+            info = json.loads(probe.stdout) if probe.returncode == 0 else {}
+        except:
+            info = {}
+        
+        duration = float(info.get("format", {}).get("duration", 0))
+        has_audio = any(s.get("codec_type") == "audio" for s in info.get("streams", []))
+        
+        # Platform-specific optimal durations
+        optimal_ranges = {
+            "tiktok": (7, 30),
+            "instagram_reels": (15, 30),
+            "youtube_shorts": (30, 60)
+        }
+        opt_min, opt_max = optimal_ranges.get(platform, (15, 60))
+        
+        # Score components
+        scores = {}
+        
+        # Length score
+        if opt_min <= duration <= opt_max:
+            scores["length_optimal"] = 90
+        elif duration < opt_min:
+            scores["length_optimal"] = max(50, 90 - (opt_min - duration) * 5)
+        else:
+            scores["length_optimal"] = max(40, 90 - (duration - opt_max) * 3)
+        
+        # Audio presence
+        scores["audio_match"] = 80 if has_audio else 30
+        
+        # Resolution check
+        for s in info.get("streams", []):
+            if s.get("codec_type") == "video":
+                h = int(s.get("height", 0))
+                w = int(s.get("width", 0))
+                if h >= 1920 or w >= 1080:
+                    scores["visual_quality"] = 85
+                elif h >= 1080:
+                    scores["visual_quality"] = 75
+                else:
+                    scores["visual_quality"] = 55
+                # Vertical check
+                if h > w:
+                    scores["format_match"] = 90
+                else:
+                    scores["format_match"] = 50
+        
+        scores.setdefault("visual_quality", 60)
+        scores.setdefault("format_match", 60)
+        scores["hook_strength"] = 70  # Can't assess without content analysis
+        scores["pacing"] = 70
+        scores["trend_alignment"] = 65
+        
+        overall = round(sum(scores.values()) / len(scores))
+        
+        suggestions = []
+        if scores.get("format_match", 0) < 70:
+            suggestions.append("Convert to 9:16 vertical format for better reach")
+        if scores.get("length_optimal", 0) < 70:
+            suggestions.append(f"Adjust length to {opt_min}-{opt_max}s for {platform}")
+        if not has_audio:
+            suggestions.append("Add audio - videos without sound get 40% less reach")
+        
+        return json.dumps({
+            "overall_score": overall,
+            "breakdown": scores,
+            "suggestions": suggestions
+        })
+
+
+class CaptionGenerator:
+    """Generate platform-optimized captions."""
+    
+    @staticmethod
+    def generate(video_description: str, platform: str, tone: str = "casual", include_cta: bool = True) -> str:
+        """Generate a caption (using the model itself for this in production)."""
+        print(f"  ✍️ Generating caption for {platform}")
+        
+        hashtag_sets = {
+            "tiktok": ["#fyp", "#viral", "#foryou", "#trending"],
+            "instagram": ["#reels", "#explore", "#instagood", "#trending"],
+            "youtube": ["#shorts", "#subscribe", "#viral"]
+        }
+        
+        base_tags = hashtag_sets.get(platform, ["#viral"])
+        
+        # Extract keywords from description for niche hashtags
+        words = video_description.lower().split()
+        niche_tags = [f"#{w}" for w in words if len(w) > 3 and w.isalpha()][:3]
+        
+        posting_times = {
+            "tiktok": "7-9am, 12-1pm, or 7-9pm in your audience timezone",
+            "instagram": "6-9am, 12-2pm, or 5-7pm EST",
+            "youtube": "2-4pm or 8-10pm EST"
+        }
+        
+        return json.dumps({
+            "caption": f"[AI will generate based on: {video_description}]",
+            "hashtags": " ".join(base_tags + niche_tags),
+            "posting_time": posting_times.get(platform, "Check your analytics"),
+            "tip": "Reply to every comment in the first hour - algorithm loves engagement"
+        })
+
+
+class AIDetector:
+    """Detect AI-generated content."""
+    
+    @staticmethod
+    def detect(content_path: str, check_type: str = "video") -> str:
+        """Basic AI content detection heuristics."""
+        print(f"  🔬 Checking for AI artifacts: {content_path}")
+        
+        if not os.path.exists(content_path):
+            return json.dumps({"error": f"File not found: {content_path}"})
+        
+        # Basic file analysis (real detection would use a classifier model)
+        size = os.path.getsize(content_path)
+        
+        return json.dumps({
+            "file_analyzed": content_path,
+            "check_type": check_type,
+            "file_size_mb": round(size / 1024 / 1024, 2),
+            "note": "Full AI detection requires DeMamba or VideoScore2 model. Basic file analysis only.",
+            "recommendations": [
+                "Check for morphing objects between frames",
+                "Look for impossible reflections or shadows",
+                "Verify text is readable and consistent",
+                "Check if camera movement is unnaturally smooth"
+            ]
+        })
+
+
+# ============================================================
+# AGENT CORE
+# ============================================================
+
+TOOL_MAP = {
+    "ffmpeg_cmd": lambda args: FFmpegTool.run(**args),
+    "web_search": lambda args: WebSearchTool.search(**args),
+    "analyze_video": lambda args: VideoAnalyzer.analyze(**args),
+    "score_virality": lambda args: ViralityScorer.score(**args),
+    "generate_caption": lambda args: CaptionGenerator.generate(**args),
+    "detect_ai_slop": lambda args: AIDetector.detect(**args),
+}
+
+
+class ViralCutAgent:
+    """The main agent that orchestrates video editing using the fine-tuned model."""
+    
+    def __init__(self, model_id="ryu34/viralcut-agent", device="auto"):
+        print(f"Loading ViralCut Agent from {model_id}...")
+        
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map=device,
+            torch_dtype="auto",
+        )
+        self.model.eval()
+        
+        # Tool definitions for the chat template
+        self.tools = [
+            {"type": "function", "function": {"name": "ffmpeg_cmd", "description": "Execute FFmpeg command for video/audio processing.", "parameters": {"type": "object", "properties": {"command": {"type": "string"}, "description": {"type": "string"}}, "required": ["command", "description"]}}},
+            {"type": "function", "function": {"name": "web_search", "description": "Search web for royalty-free assets and trends.", "parameters": {"type": "object", "properties": {"query": {"type": "string"}, "search_type": {"type": "string", "enum": ["royalty_free_music", "sound_effects", "trending_content", "general"]}}, "required": ["query", "search_type"]}}},
+            {"type": "function", "function": {"name": "analyze_video", "description": "Analyze video for scenes, audio, transcript, quality.", "parameters": {"type": "object", "properties": {"video_path": {"type": "string"}, "analysis_type": {"type": "string", "enum": ["full", "scenes", "audio", "transcript", "quality", "pacing"]}}, "required": ["video_path", "analysis_type"]}}},
+            {"type": "function", "function": {"name": "score_virality", "description": "Score video viral potential 0-100.", "parameters": {"type": "object", "properties": {"video_path": {"type": "string"}, "platform": {"type": "string", "enum": ["tiktok", "instagram_reels", "youtube_shorts"]}, "niche": {"type": "string"}}, "required": ["video_path", "platform"]}}},
+            {"type": "function", "function": {"name": "generate_caption", "description": "Generate platform-optimized caption with hashtags.", "parameters": {"type": "object", "properties": {"video_description": {"type": "string"}, "platform": {"type": "string", "enum": ["tiktok", "instagram", "youtube"]}, "tone": {"type": "string"}, "include_cta": {"type": "boolean"}}, "required": ["video_description", "platform"]}}},
+            {"type": "function", "function": {"name": "detect_ai_slop", "description": "Check content for AI-generated artifacts.", "parameters": {"type": "object", "properties": {"content_path": {"type": "string"}, "check_type": {"type": "string", "enum": ["video", "image", "text", "audio"]}}, "required": ["content_path", "check_type"]}}}
+        ]
+        
+        print("Agent ready!")
+    
+    def run(self, user_message: str, max_turns: int = 15):
+        """Run the agent on a user request, executing tool calls autonomously."""
+        
+        messages = [
+            {"role": "system", "content": "You are ViralCut Agent, an autonomous AI video editor and social media content strategist. You transform raw video footage into professional, viral-worthy social media content. Use your tools to analyze, edit, search, and optimize. Think step-by-step. Always use royalty-free content."},
+            {"role": "user", "content": user_message}
+        ]
+        
+        print(f"\n{'='*60}")
+        print(f"🎬 ViralCut Agent")
+        print(f"{'='*60}")
+        print(f"User: {user_message}\n")
+        
+        for turn in range(max_turns):
+            # Generate response
+            text = self.tokenizer.apply_chat_template(
+                messages, tools=self.tools, tokenize=False, add_generation_prompt=True
+            )
+            inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+            
+            with __import__("torch").no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=1024,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                )
+            
+            response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=False)
+            
+            # Parse response for tool calls or plain text
+            tool_calls = self._parse_tool_calls(response)
+            
+            if tool_calls:
+                # Add assistant message with tool calls
+                messages.append({"role": "assistant", "tool_calls": tool_calls})
+                
+                # Execute each tool call
+                for tc in tool_calls:
+                    func_name = tc["function"]["name"]
+                    try:
+                        args = json.loads(tc["function"]["arguments"])
+                    except:
+                        args = {}
+                    
+                    print(f"\n  🔧 Calling: {func_name}")
+                    
+                    if func_name in TOOL_MAP:
+                        result = TOOL_MAP[func_name](args)
+                    else:
+                        result = json.dumps({"error": f"Unknown tool: {func_name}"})
+                    
+                    messages.append({"role": "tool", "name": func_name, "content": result})
+                    print(f"  ✅ Result: {result[:200]}...")
+            else:
+                # Plain text response - agent is done
+                clean = self._clean_response(response)
+                messages.append({"role": "assistant", "content": clean})
+                print(f"\n🤖 Agent: {clean}")
+                break
+        
+        return messages
+    
+    def _parse_tool_calls(self, response: str) -> list:
+        """Parse tool calls from model output."""
+        tool_calls = []
+        
+        # Qwen tool call format: <tool_call>{"name": "...", "arguments": {...}}</tool_call>
+        pattern = r'<tool_call>\s*(\{.*?\})\s*</tool_call>'
+        matches = re.findall(pattern, response, re.DOTALL)
+        
+        for match in matches:
+            try:
+                data = json.loads(match)
+                tool_calls.append({
+                    "type": "function",
+                    "function": {
+                        "name": data.get("name", ""),
+                        "arguments": json.dumps(data.get("arguments", {}))
+                    }
+                })
+            except json.JSONDecodeError:
+                continue
+        
+        return tool_calls
+    
+    def _clean_response(self, response: str) -> str:
+        """Clean up model response."""
+        # Remove special tokens
+        for token in ["<|endoftext|>", "<|im_end|>", "<|im_start|>"]:
+            response = response.replace(token, "")
+        return response.strip()
+
+
+# ============================================================
+# CLI
+# ============================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="ViralCut Agent - AI Video Editor")
+    parser.add_argument("--video", type=str, help="Path to raw video file")
+    parser.add_argument("--platform", type=str, default="tiktok", 
+                       choices=["tiktok", "instagram", "youtube"],
+                       help="Target platform")
+    parser.add_argument("--niche", type=str, default="", help="Content niche")
+    parser.add_argument("--plan", action="store_true", help="Generate content plan only (no video needed)")
+    parser.add_argument("--model", type=str, default="ryu34/viralcut-agent", help="Model ID")
+    parser.add_argument("--check-slop", type=str, nargs="+", help="Check files for AI-generated content")
+    
+    args = parser.parse_args()
+    
+    if args.check_slop:
+        # Quick AI slop check without loading the full model
+        for f in args.check_slop:
+            result = AIDetector.detect(f, "video")
+            print(json.dumps(json.loads(result), indent=2))
+        return
+    
+    agent = ViralCutAgent(model_id=args.model)
+    
+    if args.plan:
+        niche = args.niche or "general"
+        agent.run(f"Research current {args.platform} trends for the '{niche}' niche and create a detailed 7-day content plan with hooks, posting times, and viral strategies.")
+    elif args.video:
+        if not os.path.exists(args.video):
+            print(f"Error: Video file not found: {args.video}")
+            sys.exit(1)
+        niche_str = f" in the {args.niche} niche" if args.niche else ""
+        agent.run(f"I have raw footage at {args.video}. Transform it into a professional, viral {args.platform} video{niche_str}. Analyze it, find the best moments, add trending music, professional edits, and optimize for maximum engagement.")
+    else:
+        # Interactive mode
+        print("ViralCut Agent - Interactive Mode")
+        print("Type your request (or 'quit' to exit):\n")
+        while True:
+            try:
+                user_input = input("You: ").strip()
+                if user_input.lower() in ("quit", "exit", "q"):
+                    break
+                if user_input:
+                    agent.run(user_input)
+            except (KeyboardInterrupt, EOFError):
+                break
+
+
+if __name__ == "__main__":
+    main()