Update app_4k.py

app_4k.py

@@ -17,6 +17,42 @@ snapshot_download(repo_id="APRIL-AIGC/UltraWan", repo_type="model", local_dir="u
 
 
 
+import subprocess
+import os
+import uuid
+
+def generate_4k_ultrawan_video(prompt: str, seed: int = 42, out_dir: str = "/tmp/output"):
+    output_name = f"ultrawan_{uuid.uuid4().hex[:8]}"
+    output_path = os.path.join(out_dir, output_name)
+    os.makedirs(out_dir, exist_ok=True)
+
+    cmd = [
+        "python", "infer.py",
+        "--model_dir", "ultrawan_weights/Wan2.1-T2V-1.3B",
+        "--model_path", "ultrawan_weights/UltraWan/ultrawan-4k.ckpt",
+        "--mode", "lora",
+        "--lora_alpha", "0.5",
+        "--usp", "0",
+        "--height", "2160",
+        "--width", "3840",
+        "--num_frames", "33",
+        "--prompt", prompt,
+        "--seed", str(seed),
+        "--out_dir", output_path
+    ]
+
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = "0"
+
+    subprocess.run(cmd, env=env, check=True)
+
+    # Assume the output is a .mp4 inside output_path
+    generated_files = [f for f in os.listdir(output_path) if f.endswith(".mp4")]
+    if not generated_files:
+        raise FileNotFoundError("No output video found.")
+
+    return os.path.join(output_path, generated_files[0])
+
 
 
 # LIGHT WEIGHT 1.3b
@@ -143,6 +179,22 @@ MAX_FRAMES_MODEL = 81
 default_prompt_t2v = "cinematic footage, group of pedestrians dancing in the streets of NYC, high quality breakdance, 4K, tiktok video, intricate details, instagram feel, dynamic camera, smooth dance motion, dimly lit, stylish, beautiful faces, smiling, music video"
 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"
 
+
+
+def load_model_from_path(model_path: str):
+    """
+    Loads a diffusion pipeline from a local directory.
+    The model is automatically loaded to CUDA with float16.
+    """
+    pipe = DiffusionPipeline.from_pretrained(
+        model_path,
+        torch_dtype=torch.float16,
+        variant="fp16" if os.path.exists(os.path.join(model_path, "model.fp16.safetensors")) else None
+    ).to("cuda")
+    pipe.enable_model_cpu_offload()  # Optional: for large models
+    return pipe
+
+    
 def get_duration(prompt, height, width, 
                    negative_prompt, duration_seconds,
                    guidance_scale, steps,
@@ -170,18 +222,10 @@ def generate_video(prompt, height, width,
 
     # Decide whether to use UltraWan or regular model
     if use_ultrawan_4k:
-        # ✅ Override with 4K resolution
-        target_h, target_w = 2160, 3840
-        steps = max(steps, 10)
-        guidance_scale = max(guidance_scale, 7.5)
-
-        # ✅ Lazy-load UltraWan model if not already loaded
-        global ultrawan_pipe
-        if "ultrawan_pipe" not in globals() or ultrawan_pipe is None:
-            from transformers import pipeline  # or appropriate loader
-            ultrawan_pipe = load_model_from_path("ultrawan_weights/UltraWan")
-
-        generator_pipe = ultrawan_pipe
+         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: