Update app_4k.py

app_4k.py

@@ -165,6 +165,32 @@ default_prompt_t2v = "cinematic footage, group of pedestrians dancing in the str
 default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature"
 
 
+import os
+import tempfile
+import random
+import numpy as np
+import torch
+import gradio as gr
+import subprocess
+import shutil
+
+def upscale_to_4k_and_replace(input_video_path):
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_upscaled:
+        upscaled_path = tmp_upscaled.name
+
+    cmd = [
+        "ffmpeg",
+        "-i", input_video_path,
+        "-vf", "scale=3840:2160:flags=lanczos",
+        "-c:v", "libx264",
+        "-crf", "18",
+        "-preset", "slow",
+        "-y",
+        upscaled_path,
+    ]
+    subprocess.run(cmd, check=True)
+
+    shutil.move(upscaled_path, input_video_path)
 
 def load_model_from_path(model_path: str):
     """
@@ -194,54 +220,53 @@ def get_duration(prompt, height, width,
 
 @spaces.GPU(duration=get_duration)
 def generate_video(prompt, height, width, 
-                   negative_prompt=default_negative_prompt, duration_seconds=2,
-                   guidance_scale=1, steps=4,
-                   seed=42, randomize_seed=False,
-                   use_ultrawan_4k=False,  # ✅ New toggle argument
+                   negative_prompt=default_negative_prompt,
+                   duration_seconds=2, guidance_scale=1,
+                   steps=4, seed=42, randomize_seed=False,
+                   use_ultrawan_4k=False,
                    progress=gr.Progress(track_tqdm=True)):
-    
+
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter a text prompt. Try to use long and precise descriptions.")
 
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
-    # Decide whether to use UltraWan or regular model
-    if use_ultrawan_4k:
-         video_path = generate_4k_ultrawan_video(prompt=prompt, seed=current_seed)
-         return video_path, current_seed
-
-        
-    else:
-        # Clamp values in demo mode
-        if IS_ORIGINAL_SPACE:
-            height = min(height, LIMITED_MAX_RESOLUTION)
-            width = min(width, LIMITED_MAX_RESOLUTION)
-            duration_seconds = min(duration_seconds, LIMITED_MAX_DURATION)
-            steps = min(steps, LIMITED_MAX_STEPS)
+    # Clamp values in demo mode
+    if IS_ORIGINAL_SPACE:
+        height = min(height, LIMITED_MAX_RESOLUTION)
+        width = min(width, LIMITED_MAX_RESOLUTION)
+        duration_seconds = min(duration_seconds, LIMITED_MAX_DURATION)
+        steps = min(steps, LIMITED_MAX_STEPS)
 
-        # Ensure height/width are valid
-        target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
-        target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
+    # Ensure dimensions are multiples of MOD_VALUE
+    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
+    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
 
-        generator_pipe = pipe  # use your existing model
+    generator_pipe = pipe  # your model pipeline
 
     num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
 
-    # Run inference
     with torch.inference_mode():
         output_frames_list = generator_pipe(
-            prompt=prompt, negative_prompt=negative_prompt,
-            height=target_h, width=target_w, num_frames=num_frames,
-            guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=target_h,
+            width=target_w,
+            num_frames=num_frames,
+            guidance_scale=float(guidance_scale),
+            num_inference_steps=int(steps),
             generator=torch.Generator(device="cuda").manual_seed(current_seed)
         ).frames[0]
 
-    # Save video
+    # Export video
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
+
     export_to_video(output_frames_list, video_path, fps=FIXED_OUTPUT_FPS)
-    upscale_to_4k(video_path, os.path.join(temp_dir, "upscaled_4k.mp4"))
 
+    # Optional 4K upscale
+    if use_ultrawan_4k:
+        upscale_to_4k_and_replace(video_path)
 
     return video_path, current_seed