Update app.py

app.py

@@ -10,10 +10,12 @@ from pathlib import Path
 from huggingface_hub import hf_hub_download
 import gradio as gr
 
-# ---------- Helper utilities ----------
+# ====================================
+# Helper utilities
+# ====================================
 
 def sh(cmd, check=True, env=None):
-    """Shell helper that streams output to stdout/stderr and returns (returncode, stdout)."""
+    """Shell helper that prints output live."""
     print(f"RUN: {cmd}")
     try:
         completed = subprocess.run(cmd, shell=True, check=check, capture_output=True, text=True, env=env)
@@ -27,235 +29,160 @@ def sh(cmd, check=True, env=None):
         print(e.stderr, file=sys.stderr)
         return e.returncode, e.stdout if hasattr(e, "stdout") else ""
 
-# ---------- FlashAttention install (best-effort) ----------
+
+# ====================================
+# FlashAttention install (startup)
+# ====================================
+
 def try_install_flash_attention():
-    """
-    Attempt to download and install the FlashAttention wheel from HF repo rahul7star/flash-attn-3
-    Path in repo: 128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl (as provided).
-    This is a best-effort install; failures are non-fatal.
-    """
-    flash_attention_installed = False
+    """Download and install FlashAttention wheel from rahul7star/flash-attn-3 repo."""
     try:
-        print("Attempting to download and install FlashAttention wheel...")
+        print("🔹 Attempting to install FlashAttention...")
         wheel = hf_hub_download(
             repo_id="rahul7star/flash-attn-3",
             repo_type="model",
             filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
         )
-        print("Downloaded wheel:", wheel)
-        rc, out = sh(f"pip install {wheel}")
-        # refresh site-packages so Python can see newly-installed extension
-        try:
-            import importlib, site
-            # add site-packages dir (first one) and invalidate caches
-            site.addsitedir(site.getsitepackages()[0])
-            importlib.invalidate_caches()
-        except Exception as e:
-            print("Could not update site-packages cache:", e)
-        flash_attention_installed = True
-        print("FlashAttention installed successfully.")
+        print(f"✅ Wheel downloaded: {wheel}")
+        sh(f"pip install {wheel}")
+        import importlib, site
+        site.addsitedir(site.getsitepackages()[0])
+        importlib.invalidate_caches()
+        print("✅ FlashAttention installed successfully.")
     except Exception as e:
         print(f"⚠️ Could not install FlashAttention: {e}")
-        print("Continuing without FlashAttention...")
-    return flash_attention_installed
-
-# ---------- Model downloader ----------
-def ensure_models_downloaded(marker_file=".models_ready"):
-    """
-    Run download_models.py if models haven't been downloaded yet.
-    This creates a small marker file after success to avoid repeated downloads.
-    """
-    marker = Path(marker_file)
+        print("Continuing without it...")
+
+
+# ====================================
+# Model download (startup)
+# ====================================
+
+def ensure_models_downloaded():
+    """Run download_models.py once at startup to fetch model weights."""
+    marker = Path(".models_ready")
     if marker.exists():
-        print("Models already downloaded (marker found).")
+        print("✅ Models already downloaded (marker found).")
         return True
 
     if not Path("download_models.py").exists():
-        print("Warning: download_models.py not found in repo. Please add it or run model download manually.")
+        print("❌ Missing download_models.py in repo. Please include it.")
         return False
 
+    print("⬇️ Downloading model weights via download_models.py ...")
     try:
-        print("Running download_models.py to fetch model artifacts...")
-        # Try to call the script directly. Use same python executable.
-        rc, out = sh(f"{sys.executable} download_models.py", check=True)
-        # If it completes without exception, create marker
+        rc, _ = sh(f"{sys.executable} download_models.py", check=True)
         marker.write_text("ok")
-        print("download_models.py finished. Marker created.")
+        print("✅ Model download complete.")
         return True
     except Exception as e:
-        print("Failed to run download_models.py:", e)
+        print(f"❌ Model download failed: {e}")
         return False
 
-# ---------- Inference runner ----------
-def run_inference(prompt: str, image_path: str | None, seed: int | None = None, duration: float | None = None, workdir: str | None = None):
-    """
-    Run test.py with prompt and optional image. Expect test.py to produce a video file (e.g. output.mp4)
-    Returns path to produced video or None on failure.
-    """
-    workdir = workdir or os.getcwd()
-    out_video = Path(workdir) / "output.mp4"
 
-    # remove old output if present
-    if out_video.exists():
-        try:
-            out_video.unlink()
-        except Exception:
-            pass
+# ====================================
+# Inference runner (text/image → video)
+# ====================================
 
-    if not Path("test.py").exists():
-        raise FileNotFoundError("test.py not found in repo. Place the repo's test.py in the same folder as app.py.")
+def run_inference(prompt: str, image_path: str | None = None):
+    """Run test.py with prompt + optional image. Returns path to video."""
+    workdir = os.getcwd()
+    out_video = Path(workdir) / "output.mp4"
+    if out_video.exists():
+        out_video.unlink(missing_ok=True)
 
     cmd = [sys.executable, "test.py", "--prompt", f"\"{prompt}\""]
     if image_path:
         cmd += ["--image_path", f"\"{image_path}\""]
-    if seed is not None:
-        cmd += ["--seed", str(seed)]
-    if duration is not None:
-        # If the test.py uses a --duration flag; adapt if your script uses different arg name.
-        cmd += ["--duration", str(duration)]
 
-    # Join to single command string to ensure shell wildcard expansion if needed
     cmd_str = " ".join(cmd)
-    print("Inference command:", cmd_str)
+    print(f"🚀 Running inference: {cmd_str}")
 
     try:
-        # We stream output and check for completion
-        proc = subprocess.run(cmd_str, shell=True, check=True, capture_output=True, text=True, env=os.environ)
-        print("Inference stdout:", proc.stdout)
+        proc = subprocess.run(cmd_str, shell=True, capture_output=True, text=True, check=True)
+        print(proc.stdout)
         if proc.stderr:
-            print("Inference stderr:", proc.stderr, file=sys.stderr)
+            print(proc.stderr, file=sys.stderr)
     except subprocess.CalledProcessError as e:
-        print("Inference failed:", e, file=sys.stderr)
-        print(e.stdout if hasattr(e, "stdout") else "")
-        print(e.stderr if hasattr(e, "stderr") else "", file=sys.stderr)
+        print("❌ Inference failed:", e)
+        print(e.stdout)
+        print(e.stderr)
         return None
 
-    # locate output video
+    # Find the resulting .mp4
     if out_video.exists():
         return str(out_video)
-    # fallback: find any recent mp4 in workdir
-    candidates = sorted(Path(workdir).glob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True)
-    if candidates:
-        return str(candidates[0])
-    return None
-
-# ---------- Gradio app callbacks ----------
-@spaces.GPU(duration = 50)
-def generate(prompt, image, seed, duration, install_flash, force_download_models):
-    """
-    Main callback for Gradio "Generate" button.
-    - install_flash: boolean, whether to attempt flash-attn install this run
-    - force_download_models: boolean to re-run download_models.py even if marker exists
-    Returns (video_file, status_text)
-    """
-    status_msgs = []
-    # Convert image (gradio gives a PIL Image or None) to a temp file if provided
-    temp_image_path = None
+    vids = sorted(Path(workdir).glob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True)
+    return str(vids[0]) if vids else None
+
+
+# ====================================
+# Gradio callback
+# ====================================
+
+@spaces.GPU(duration=50)
+def generate(prompt, image):
+    """Main Gradio callback for generating video."""
+    status = []
+    temp_img_path = None
+
     if image is not None:
         tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        try:
-            image.save(tmp, format="PNG")
-            tmp.flush()
-            temp_image_path = tmp.name
-            tmp.close()
-            status_msgs.append(f"Saved input image to {temp_image_path}")
-        except Exception as e:
-            status_msgs.append(f"Failed to save uploaded image: {e}")
-            temp_image_path = None
-
-    # Optionally install flash attention
-    if install_flash:
-        ok = try_install_flash_attention()
-        status_msgs.append(f"Attempted FlashAttention install: {'OK' if ok else 'FAILED'}")
-    else:
-        status_msgs.append("Skipped FlashAttention install (checkbox unchecked).")
-
-    # Ensure models downloaded
-    if force_download_models:
-        # remove marker if present so we re-download
-        marker = Path(".models_ready")
-        if marker.exists():
-            try:
-                marker.unlink()
-                status_msgs.append("Removed existing model marker to force re-download.")
-            except Exception as e:
-                status_msgs.append(f"Could not remove marker file: {e}")
-
-    ok_models = ensure_models_downloaded()
-    status_msgs.append(f"Models ready: {'yes' if ok_models else 'no'}")
-    if not ok_models:
-        status_msgs.append("Warning: models not ready. Inference will probably fail.")
-
-    # Run inference
-    status_msgs.append("Starting inference (this may take time on GPU).")
+        image.save(tmp, format="PNG")
+        tmp.close()
+        temp_img_path = tmp.name
+        status.append(f"📸 Saved image: {temp_img_path}")
+
     try:
-        video_path = run_inference(prompt=prompt, image_path=temp_image_path, seed=seed, duration=duration)
+        video_path = run_inference(prompt, image_path=temp_img_path)
+        if not video_path:
+            status.append("❌ No video produced. Check test.py output.")
+            return None, "\n".join(status)
     except Exception as e:
-        status_msgs.append(f"Inference runner raised an exception: {e}")
-        return None, "\n".join(status_msgs)
-
-    if video_path:
-        status_msgs.append(f"Video created: {video_path}")
-        # Move to /tmp or keep in repo for Gradio to serve
-        # We'll copy to a stable path that Gradio can serve e.g. ./outputs/output_{timestamp}.mp4
-        dest_dir = Path("outputs")
-        dest_dir.mkdir(exist_ok=True)
-        ts = int(time.time())
-        dest = dest_dir / f"t2v_output_{ts}.mp4"
-        try:
-            shutil.copy(video_path, dest)
-            status_msgs.append(f"Video copied to {dest}")
-            return str(dest), "\n".join(status_msgs)
-        except Exception as e:
-            status_msgs.append(f"Could not copy video to outputs/: {e}")
-            # still try to return original path
-            return str(video_path), "\n".join(status_msgs)
-    else:
-        status_msgs.append("No video produced by test.py (output not found). Check logs.")
-        return None, "\n".join(status_msgs)
-
-# ---------- Build Gradio interface ----------
-def build_ui():
-    with gr.Blocks(title="Text+Image → Video (Spaces GPU)", css="""
-        .output-video { max-width: 800px; }
-    """) as demo:
-        gr.Markdown("# Text + (Optional) Image → Video\nSimple UI to run Kandinsky/Wan T2V `test.py` in this Space (GPU required).")
+        status.append(f"❌ Inference failed: {e}")
+        return None, "\n".join(status)
+
+    dest_dir = Path("outputs"); dest_dir.mkdir(exist_ok=True)
+    ts = int(time.time())
+    dest = dest_dir / f"t2v_output_{ts}.mp4"
+    shutil.copy(video_path, dest)
+    status.append(f"✅ Video generated: {dest}")
+    return str(dest), "\n".join(status)
 
+
+# ====================================
+# UI builder
+# ====================================
+
+def build_ui():
+    with gr.Blocks(title="Text+Image → Video (Spaces GPU)") as demo:
+        gr.Markdown("## 🎬 Kandinsky / T2V Video Generator\nProvide a text prompt and optional image to generate short video clips using GPU inference.")
         with gr.Row():
             with gr.Column(scale=3):
-                prompt = gr.Textbox(label="Prompt", placeholder="A dog in a red hat, cinematic, 5s", value="A dog in a red hat")
-                image_in = gr.Image(label="Optional reference image (still)", type="pil")
-                with gr.Row():
-                    seed = gr.Number(value=42, label="Seed (optional)", precision=0)
-                    duration = gr.Number(value=5.0, label="Duration (seconds, optional)", precision=2)
-                install_flash = gr.Checkbox(label="Attempt FlashAttention install before running (best-effort)", value=False)
-                force_download = gr.Checkbox(label="Force run download_models.py (re-download models)", value=False)
-                generate_btn = gr.Button("Generate Video", variant="primary")
-                status = gr.Textbox(label="Status / Logs", interactive=False, lines=10)
+                prompt = gr.Textbox(label="Prompt", placeholder="A dog in a red hat, cinematic lighting", value="A dog in a red hat")
+                image_in = gr.Image(label="Optional input image", type="pil")
+                generate_btn = gr.Button("🎥 Generate Video", variant="primary")
+                status = gr.Textbox(label="Logs", lines=8)
             with gr.Column(scale=2):
-                out_video = gr.Video(label="Output video", elem_classes="output-video")
-                gr.Markdown("**Notes**:\n- Ensure `download_models.py` and `test.py` are present and compatible.\n- `test.py` should produce an mp4 named `output.mp4` in the repo root or an mp4 somewhere in the working dir.\n- Long-running jobs may hit Space runtime limits if very long.")
-
-        # wire up
-        generate_btn.click(fn=generate,
-                           inputs=[prompt, image_in, seed, duration, install_flash, force_download],
-                           outputs=[out_video, status])
+                out_video = gr.Video(label="Output video")
 
+        generate_btn.click(fn=generate, inputs=[prompt, image_in], outputs=[out_video, status])
     return demo
 
-# ---------- Main entrypoint ----------
+
+# ====================================
+# App startup
+# ====================================
+
 if __name__ == "__main__":
-    # Quick environment checks
-    print("Starting T2V Gradio app. Python:", sys.executable)
-    print("CUDA available?", os.environ.get("CUDA_VISIBLE_DEVICES", "(not set)"))
-    # Attempt to install flash-attn automatically? We default to not attempting until user requests in UI.
-    # Pre-check models: create marker if download_models.py has already run previously
-    if not Path(".models_ready").exists() and Path("download_models.py").exists():
-        # we do NOT force downloading on startup automatically to avoid long startup delays on Spaces.
-        print("download_models.py exists. Models not yet marked as downloaded. Use the UI to run download (or set force flag).")
-
-    # Create outputs dir
-    Path("outputs").mkdir(exist_ok=True)
+    print("🚀 Starting Text+Image → Video Gradio App")
+    print("Python:", sys.executable)
+    print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES", "(not set)"))
+
+    # Install FlashAttention + download models ONCE at startup
+    try_install_flash_attention()
+    ensure_models_downloaded()
 
+    Path("outputs").mkdir(exist_ok=True)
     demo = build_ui()
     demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))