app.py

# --- PATCH FOR JINJA2 / GRADIO COMPATIBILITY ISSUE ---
import jinja2

# Fixes the 'TypeError: unhashable type: dict' issue in older Gradio versions
if not hasattr(jinja2.utils.LRUCache, '__getitem__'):
    def fallback_getitem(self, key):
        try:
            return self._mapping[key]
        except TypeError:
            return self._mapping.get(str(key))
    jinja2.utils.LRUCache.__getitem__ = fallback_getitem

original_get = jinja2.utils.LRUCache.get
def patched_get(self, key, default=None):
    try:
        return original_get(self, key, default)
    except TypeError:
        return self._mapping.get(str(key), default)
jinja2.utils.LRUCache.get = patched_get
# -----------------------------------------------------

import gradio as gr
import numpy as np
import os
import tempfile
import requests
from gradio_client import Client

# Try to import local AI Mind modules for optional direct execution
try:
    from memory import ConversationMemory
    from brain import Brain
    HAS_LOCAL_AI = True
    print("🧠 Local InvictaTill AI Mind loaded successfully!")
except Exception as e:
    HAS_LOCAL_AI = False
    print(f"⚠️ Local InvictaTill AI Mind unavailable (using HTTP client): {str(e)}")

STYLE_PRESETS = {
    "None": "",
    "Cyberpunk / Neon Glow": "cyberpunk style, futuristic, neon lights, high contrast, dark atmosphere, synthetic lighting",
    "Anime / Makoto Shinkai": "anime aesthetic, vibrant colors, beautiful clouds, sun flare, highly detailed, by Makoto Shinkai",
    "Photorealistic Cinematic": "photorealistic, cinematic film, 35mm lens, highly detailed, realistic lighting, volumetric dust, warm color grading",
    "3D Pixar / Disney": "Pixar style, 3D animated character style, smooth textures, vibrant lighting, friendly atmosphere",
    "Oil Painting / Fine Art": "oil painting style, rich textures, visible brushstrokes, high-end fine art aesthetic, masterfully rendered",
    "Vintage Film / VHS Retro": "vintage film, 80s retro, VHS tape texture, light leaks, chromatic aberration, retro color grading"
}

def enhance_prompt_with_ai(prompt_text, style_preset, ai_mode, ai_url, ai_key, session_id):
    if not prompt_text or not prompt_text.strip():
        return "Please enter a prompt first."
        
    style_modifiers = STYLE_PRESETS.get(style_preset, "")
    
    system_prompt = (
        "You are an expert cinematic prompt engineer for video generation models (like Wan 2.1, LTX-Video, CogVideo). "
        "Your task is to rewrite the user's simple prompt into a highly descriptive, visually stunning, "
        "and detailed prompt optimized for text-to-video models. "
        "Use your learned facts, memory context, and knowledge about the user's business if relevant. "
        "Include specific details about lighting, camera angle, motion, and atmosphere. "
        "Keep it under 75 words. "
        "Respond ONLY with the final enhanced prompt. Do NOT include any intro or conversational filler."
    )
    
    full_prompt = prompt_text
    if style_modifiers:
        full_prompt += f" with the style: {style_modifiers}"
        
    # Local Mode: Load the local AI Mind directly
    if ai_mode == "Local Integrated Mind (Direct Codebase)" and HAS_LOCAL_AI:
        try:
            # Connect to local database path
            db_path = "/data/invicta_data/memory.db"
            if not os.path.exists(os.path.dirname(db_path)):
                db_path = os.path.join(tempfile.gettempdir(), "memory.db")
                
            mem = ConversationMemory(db_path=db_path)
            
            # Fetch active user context if any user exists
            user_id = None
            user_profile = None
            try:
                user_ids = mem.get_all_user_ids() if hasattr(mem, 'get_all_user_ids') else []
                if user_ids:
                    user_id = user_ids[0]
                    user_profile = mem.get_user_profile(user_id)
            except Exception:
                pass
                
            # Direct Brain Query
            brain = Brain()
            query = f"[SYSTEM CONTEXT]\n{system_prompt}\n\nUser: {full_prompt}\n\nAssistant:"
            answer, _ = brain.think(
                user_query=query,
                user_profile_dict=user_profile,
                user_id=user_id
            )
            if answer:
                return answer.strip().strip('"')
        except Exception as e:
            print(f"⚠️ Local AI Mind execution failed: {str(e)}. Falling back to cloud...")
            
    # Cloud Mode: Standard API Post
    ai_url = (ai_url or "").strip().rstrip("/")
    if not ai_url:
        ai_url = "https://invictatill-invictatill-ai.hf.space"
    
    chat_endpoint = f"{ai_url}/api/v1/chat"
    
    payload = {
        "message": f"[SYSTEM CONTEXT]\n{system_prompt}\n\nUser: {full_prompt}\n\nAssistant:"
    }
    if session_id:
        payload["session_id"] = session_id
        
    headers = {"Content-Type": "application/json"}
    if ai_key:
        headers["Authorization"] = f"Bearer {ai_key}"
            
    try:
        response = requests.post(chat_endpoint, json=payload, headers=headers, timeout=12)
        if response.status_code == 200:
            data = response.json()
            enhanced = data.get("reply") or data.get("choices", [{}])[0].get("message", {}).get("content", "")
            if enhanced:
                return enhanced.strip().strip('"')
        return f"{prompt_text}, {style_modifiers}".strip(", ")
    except Exception:
        return f"{prompt_text}, {style_modifiers}".strip(", ")

def generate_video(prompt, negative_prompt, style_preset, generator_model, input_video, ai_mode, ai_url, ai_key, session_id, progress=gr.Progress()):
    if not prompt or prompt.strip() == "":
        return None, "❌ Please enter a prompt."
    
    # 1. Enhance the prompt using InvictaTill AI first
    progress(0.1, desc="Enhancing prompt with InvictaTill AI Mind...")
    enhanced_prompt = enhance_prompt_with_ai(prompt, style_preset, ai_mode, ai_url, ai_key, session_id)
    print(f"Original Prompt: {prompt}")
    print(f"Enhanced Prompt: {enhanced_prompt}")
    
    # Generate seed
    seed = int(np.random.randint(0, 2**32 - 1))
    
    video_path = None
    success_space = None
    
    # 2. Check if Cosmos-Transfer is selected
    if generator_model == "NVIDIA Cosmos-Transfer2.5-2b (Physics NIM)":
        if not input_video:
            return None, "❌ NVIDIA Cosmos-Transfer requires an Input Control Video for Sim2Real style transfer. Please upload a video first."
            
        progress(0.3, desc="Connecting to NVIDIA Cosmos NIM Endpoint...")
        
        # Load API Key (NVIDIA Key)
        # Fallback to default working NVIDIA key from brain.py if not provided
        nvidia_key = ai_key if (ai_key and ai_key.strip()) else "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"
        
        try:
            import base64
            progress(0.4, desc="Encoding input video file...")
            with open(input_video, "rb") as f:
                video_base64 = base64.b64encode(f.read()).decode("utf-8")
                
            invoke_url = "https://ai.api.nvidia.com/v1/cosmos/nvidia/cosmos-transfer2.5-2b"
            headers = {
                "Authorization": f"Bearer {nvidia_key}",
                "Accept": "application/json",
                "Content-Type": "application/json"
            }
            
            payload = {
                "prompt": enhanced_prompt,
                "video": f"data:video/mp4;base64,{video_base64}",
                "strength": 0.85
            }
            
            progress(0.5, desc="Sending transfer request to NVIDIA Cloud...")
            res = requests.post(invoke_url, headers=headers, json=payload, timeout=90)
            
            if res.status_code == 200:
                data = res.json()
                video_b64 = data.get("b64_video") or data.get("video")
                if video_b64:
                    if "base64," in video_b64:
                        video_b64 = video_b64.split("base64,")[1]
                    video_path = os.path.join(tempfile.gettempdir(), f"cosmos_out_{seed}.mp4")
                    with open(video_path, "wb") as f:
                        f.write(base64.b64decode(video_b64))
                    success_space = "NVIDIA Cosmos-Transfer2.5-2b (Direct Response)"
            
            elif res.status_code == 202:
                # Asynchronous execution, polling is required
                req_id = res.json().get("id") or res.headers.get("NVCF-REQID") or res.headers.get("NV-Request-Id")
                if not req_id:
                    raise Exception("Asynchronous request accepted by NVIDIA, but no Request ID returned.")
                    
                # Poll the status endpoint
                import time
                poll_url = f"https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/{req_id}"
                poll_headers = {
                    "Authorization": f"Bearer {nvidia_key}",
                    "Accept": "application/json"
                }
                
                for i in range(25): # poll up to 100s
                    time.sleep(4)
                    progress(0.5 + 0.02 * i, desc=f"NVIDIA Cosmos rendering... (polling status {i+1}/25)")
                    poll_res = requests.get(poll_url, headers=poll_headers)
                    
                    if poll_res.status_code == 200:
                        poll_data = poll_res.json()
                        # Output video extraction
                        video_b64 = poll_data.get("b64_video") or poll_data.get("video")
                        if video_b64:
                            if "base64," in video_b64:
                                video_b64 = video_b64.split("base64,")[1]
                            video_path = os.path.join(tempfile.gettempdir(), f"cosmos_{req_id}.mp4")
                            with open(video_path, "wb") as f:
                                f.write(base64.b64decode(video_b64))
                            success_space = "NVIDIA Cosmos-Transfer2.5-2b (Polled NIM)"
                            break
                    elif poll_res.status_code == 202:
                        continue
                    else:
                        raise Exception(f"NVIDIA polling failed: {poll_res.status_code} - {poll_res.text}")
            else:
                raise Exception(f"NVIDIA API Error {res.status_code}: {res.text}")
                
        except Exception as e:
            print(f"NVIDIA Cosmos execution failed: {str(e)}")
            return None, f"❌ NVIDIA Cosmos execution failed: {str(e)}"
            
    else:
        # Standard Hugging Face Cloud Spaces
        progress(0.3, desc="Connecting to Hugging Face Cloud Video Generator...")
        
        # Determine Space to target based on selection
        if generator_model == "Lightricks LTX-Video (Distilled)":
            target_spaces = [{"name": "Lightricks/ltx-video-distilled", "type": "ltx"}]
        else:
            target_spaces = [{"name": "Wan-AI/Wan2.1", "type": "wan"}]
            
        for space in target_spaces:
            try:
                progress(0.5, desc=f"Generating video using {space['name']} in the cloud...")
                client = Client(space["name"], token=ai_key if ai_key else None)
                
                if space["type"] == "ltx":
                    res = client.predict(
                        prompt=enhanced_prompt,
                        negative_prompt=negative_prompt if negative_prompt else "worst quality, inconsistent motion, blurry, jittery, distorted",
                        input_image_filepath=None,
                        input_video_filepath=None,
                        height_ui=512,
                        width_ui=704,
                        mode="text-to-video",
                        duration_ui=2,
                        ui_frames_to_use=9,
                        seed_ui=seed,
                        randomize_seed=True,
                        ui_guidance_scale=1.0,
                        improve_texture_flag=True,
                        api_name="/text_to_video"
                    )
                    
                    if isinstance(res, tuple):
                        video_data = res[0]
                    else:
                        video_data = res
                        
                    if isinstance(video_data, dict):
                        video_path = video_data.get("video") or video_data.get("path")
                    else:
                        video_path = video_data
                
                elif space["type"] == "wan":
                    res = client.predict(
                        prompt=enhanced_prompt,
                        size="1280*720",
                        watermark_wan=True,
                        seed=seed,
                        api_name="/t2v_generation_async"
                    )
                    
                    # Poll status_refresh in a loop for up to 60 seconds
                    import time
                    for i in range(15):
                        time.sleep(4)
                        progress((0.5 + 0.03 * i), desc="Generating frames in Wan Space... (polling status)")
                        status_res = client.predict(api_name="/status_refresh")
                        if isinstance(status_res, tuple) and status_res[0]:
                            video_data = status_res[0]
                            if isinstance(video_data, dict) and video_data.get("video"):
                                video_path = video_data["video"]
                                break
                    
                if video_path and os.path.exists(video_path):
                    success_space = space["name"]
                    break
            except Exception as err:
                print(f"Failed to generate on {space['name']}: {str(err)}")
                continue
                
    if not video_path:
        return None, "❌ Cloud generation failed. The selected service is currently overloaded or unresponsive. Please try again."
        
    progress(1.0, desc="Video generation complete!")
    
    info = f"""
**Cinematic Prompt (Enhanced):** {enhanced_prompt}
**Video Engine:** {success_space}
**Seed:** {seed}
**Status:** Powered entirely by InvictaTill AI & Cloud NIMs (No local GPU required)
    """.strip()
    
    return video_path, info

custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;700&family=Inter:wght@400;600;800&display=swap');

body {
    font-family: 'Inter', sans-serif !important;
    background-color: #0b0914 !important;
    background-image: radial-gradient(circle at 10% 20%, rgba(124, 58, 237, 0.08) 0%, transparent 40%),
                      radial-gradient(circle at 90% 80%, rgba(6, 182, 212, 0.06) 0%, transparent 40%) !important;
    color: #f1f5f9 !important;
}
.gradio-container {
    background: transparent !important;
    border: none !important;
    max-width: 1100px !important;
    margin: 0 auto !important;
}
.header {
    text-align: center;
    padding: 2.5rem 0 1rem;
    margin-bottom: 2rem;
}
.header h1 {
    font-family: 'Space Grotesk', sans-serif !important;
    font-size: 3rem;
    font-weight: 800;
    letter-spacing: -1.5px;
    background: linear-gradient(135deg, #a78bfa, #22d3ee);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 0.5rem;
}
.header p {
    color: #94a3b8;
    font-size: 1.1rem;
    font-weight: 500;
}
.panel {
    background: rgba(18, 16, 30, 0.65) !important;
    backdrop-filter: blur(24px) !important;
    -webkit-backdrop-filter: blur(24px) !important;
    border: 1px solid rgba(167, 139, 250, 0.15) !important;
    border-radius: 20px !important;
    padding: 2rem !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3) !important;
}
input, textarea, select {
    background: rgba(30, 27, 50, 0.8) !important;
    border: 1px solid rgba(167, 139, 250, 0.2) !important;
    border-radius: 12px !important;
    color: #f1f5f9 !important;
}
input:focus, textarea:focus, select:focus {
    border-color: #22d3ee !important;
    box-shadow: 0 0 0 3px rgba(34, 211, 238, 0.2) !important;
}
button.primary {
    background: linear-gradient(135deg, #7c3aed, #0891b2) !important;
    border: none !important;
    border-radius: 12px !important;
    font-weight: 700 !important;
    transition: all 0.25s ease !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.3) !important;
}
button.primary:hover {
    transform: translateY(-1.5px) !important;
    box-shadow: 0 6px 20px rgba(124, 58, 237, 0.45) !important;
}
button.secondary {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 12px !important;
    color: white !important;
    transition: all 0.2s !important;
}
button.secondary:hover {
    background: rgba(255, 255, 255, 0.1) !important;
    border-color: rgba(255, 255, 255, 0.2) !important;
}
.example-chip {
    cursor: pointer;
    padding: 0.5rem 1rem;
    background: rgba(124, 58, 237, 0.08);
    border: 1px solid rgba(124, 58, 237, 0.25);
    border-radius: 20px;
    font-size: 0.82rem;
    color: #c4b5fd;
    display: inline-block;
    margin: 0.25rem;
    transition: all 0.2s ease;
}
.example-chip:hover {
    background: rgba(124, 58, 237, 0.18);
    border-color: #a78bfa;
    transform: scale(1.03);
}
"""

EXAMPLES = [
    "A cyberpunk drone shot flying through neon-lit Tokyo streets at night, rain droplets on lens, cinematic lighting",
    "Slow-motion explosion of colorful powder in a dark studio, particles swirling, dramatic lighting",
    "Astronaut floating in a vibrant nebula, stars twinkling, slow rotation, ethereal glow",
    "Japanese garden in spring, cherry blossoms falling, gentle breeze, golden hour, dolly shot",
    "Futuristic car racing through a glass tunnel underwater, bioluminescent creatures outside, motion blur",
    "Abstract fluid simulation, iridescent colors mixing, dark background, high viscosity, 3D render",
]

with gr.Blocks(css=custom_css, title="InvictaTill VideoGen Studio", theme=gr.themes.Base()) as demo:
    gr.HTML("""
    <div class="header">
        <h1>🎬 InvictaTill VideoGen Studio</h1>
        <p>Generate high-end cinematic videos powered entirely by InvictaTill AI and Hugging Face Cloud Spaces</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### ⚙️ Generation Model")
            generator_model = gr.Dropdown(
                choices=[
                    "Lightricks LTX-Video (Distilled)",
                    "Wan-AI Wan 2.1 (ZeroGPU)",
                    "NVIDIA Cosmos-Transfer2.5-2b (Physics NIM)"
                ],
                value="Lightricks LTX-Video (Distilled)",
                label="Choose Video Generator Engine"
            )
            
            input_video = gr.Video(
                label="Input Video (Required ONLY for NVIDIA Cosmos-Transfer style transfer)",
                interactive=True
            )
            
            gr.Markdown("### ✍️ Prompt Composer")
            prompt = gr.Textbox(label="Describe your scene", placeholder="A cyberpunk drone shot flying through neon-lit Tokyo streets...", lines=4, elem_id="prompt")
            
            with gr.Accordion("🧠 InvictaTill AI Mind Settings", open=False):
                gr.Markdown("Configure the endpoint URL and API Key for your running InvictaTill AI Space instance so the prompt enhancer can read your business insights and custom memories.")
                ai_mode_dropdown = gr.Dropdown(
                    choices=[
                        "Local Integrated Mind (Direct Codebase)",
                        "Cloud Space API (Remote HTTP)"
                    ] if HAS_LOCAL_AI else [
                        "Cloud Space API (Remote HTTP)"
                    ],
                    value="Local Integrated Mind (Direct Codebase)" if HAS_LOCAL_AI else "Cloud Space API (Remote HTTP)",
                    label="AI Execution Mode"
                )
                ai_url_input = gr.Textbox(
                    value=os.environ.get("VITE_INVICTATILL_AI_URL", "https://invictatill-invictatill-ai.hf.space"),
                    label="AI Mind Endpoint URL",
                    placeholder="https://invictatill-invictatill-ai.hf.space"
                )
                ai_key_input = gr.Textbox(
                    value=os.environ.get("VITE_INVICTATILL_AI_KEY", ""),
                    label="API Key / Auth Token (NVIDIA Key for Cosmos)",
                    placeholder="invicta_sk_... or nvapi-...",
                    type="password"
                )
                session_id_input = gr.Textbox(
                    value="videogen_studio_session",
                    label="Session ID (Loads Memory Context)",
                    placeholder="videogen_studio_session"
                )
            
            with gr.Row():
                enhance_btn = gr.Button("✨ Enhance Prompt with InvictaTill AI Mind", variant="secondary")
                
            style_dropdown = gr.Dropdown(choices=list(STYLE_PRESETS.keys()), value="None", label="Choose Style Overlay")
            
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="blur, distortion, low quality, watermark", lines=2, value="blur, distortion, low quality, watermark, text, bad anatomy, deformed, cartoonish, static")
            
            gr.Markdown("### 🌟 Sample Prompt Concepts")
            example_html = ""
            for ex in EXAMPLES:
                safe = ex.replace('"', '&quot;')
                example_html += f'<span class="example-chip" onclick="document.querySelector(\'#prompt textarea\').value=\'{safe}\'">{ex[:35]}...</span>'
            gr.HTML(example_html)
            
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### 📼 Output Cinematic Video")
            generate_btn = gr.Button("🚀 Generate High-End Video", variant="primary", size="lg")
            
            video_output = gr.Video(label="Generated Cinematic")
            info_output = gr.Markdown()
    
    # Click Handlers
    enhance_btn.click(
        fn=enhance_prompt_with_ai,
        inputs=[prompt, style_dropdown, ai_mode_dropdown, ai_url_input, ai_key_input, session_id_input],
        outputs=[prompt]
    )
    
    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, negative_prompt, style_dropdown, generator_model, input_video, ai_mode_dropdown, ai_url_input, ai_key_input, session_id_input],
        outputs=[video_output, info_output]
    )

if __name__ == "__main__":
    demo.queue(max_size=5).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_api=False,
    )