feat: implement NVIDIA Cosmos-Transfer2.5-2b video-to-video style transfer support with base64 video encoding and NVCF status polling

app.py

@@ -128,7 +128,7 @@ def enhance_prompt_with_ai(prompt_text, style_preset, ai_mode, ai_url, ai_key, s
     except Exception:
         return f"{prompt_text}, {style_modifiers}".strip(", ")
 
-def generate_video(prompt, negative_prompt, style_preset, ai_mode, ai_url, ai_key, session_id, progress=gr.Progress()):
+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."
     
@@ -141,91 +141,176 @@ def generate_video(prompt, negative_prompt, style_preset, ai_mode, ai_url, ai_ke
     # Generate seed
     seed = int(np.random.randint(0, 2**32 - 1))
     
-    # 2. Call the Cloud Space API to generate the video
-    progress(0.3, desc="Connecting to Hugging Face Cloud Video Generator...")
-    
-    # List of verified public spaces to try sequentially (fault-tolerance)
-    spaces = [
-        {"name": "Lightricks/ltx-video-distilled", "type": "ltx"},
-        {"name": "Wan-AI/Wan2.1", "type": "wan"}
-    ]
-    
     video_path = None
     success_space = None
     
-    for space in 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)
+    # 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."
             
-            if space["type"] == "ltx":
-                # Predict signature: predict(prompt, negative_prompt, input_image_filepath, input_video_filepath, height_ui, width_ui, mode, duration_ui, ui_frames_to_use, seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag, api_name="/text_to_video")
-                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"
-                )
+        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")
                 
-                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
+            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"
+            }
             
-            elif space["type"] == "wan":
-                # Predict signature: predict(prompt, size, watermark_wan, seed, api_name="/t2v_generation_async")
-                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
+            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
-                for i in range(15):
+                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.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"]
+                    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}")
                 
-            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
+        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. All public spaces are currently overloaded. Please try again in a few moments."
+        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}
-**Cloud Model:** {success_space}
+**Video Engine:** {success_space}
 **Seed:** {seed}
-**Status:** Powered entirely by InvictaTill AI & HF Cloud Spaces (No local GPU required)
+**Status:** Powered entirely by InvictaTill AI & Cloud NIMs (No local GPU required)
     """.strip()
     
     return video_path, info
@@ -347,6 +432,22 @@ with gr.Blocks(css=custom_css, title="InvictaTill VideoGen Studio", theme=gr.the
     
     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")
             
@@ -369,8 +470,8 @@ with gr.Blocks(css=custom_css, title="InvictaTill VideoGen Studio", theme=gr.the
                 )
                 ai_key_input = gr.Textbox(
                     value=os.environ.get("VITE_INVICTATILL_AI_KEY", ""),
-                    label="API Key / Auth Token",
-                    placeholder="invicta_sk_...",
+                    label="API Key / Auth Token (NVIDIA Key for Cosmos)",
+                    placeholder="invicta_sk_... or nvapi-...",
                     type="password"
                 )
                 session_id_input = gr.Textbox(
@@ -409,7 +510,7 @@ with gr.Blocks(css=custom_css, title="InvictaTill VideoGen Studio", theme=gr.the
     
     generate_btn.click(
         fn=generate_video,
-        inputs=[prompt, negative_prompt, style_dropdown, ai_mode_dropdown, ai_url_input, ai_key_input, session_id_input],
+        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]
     )