Create app.py

app.py

@@ -0,0 +1,260 @@
+# app.py
+import os
+import sys
+import time
+import subprocess
+import tempfile
+import shutil
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+import gradio as gr
+
+# ---------- Helper utilities ----------
+
+def sh(cmd, check=True, env=None):
+    """Shell helper that streams output to stdout/stderr and returns (returncode, stdout)."""
+    print(f"RUN: {cmd}")
+    try:
+        completed = subprocess.run(cmd, shell=True, check=check, capture_output=True, text=True, env=env)
+        print(completed.stdout)
+        if completed.stderr:
+            print("ERR:", completed.stderr, file=sys.stderr)
+        return completed.returncode, completed.stdout
+    except subprocess.CalledProcessError as e:
+        print("Command failed:", e, file=sys.stderr)
+        print(e.stdout)
+        print(e.stderr, file=sys.stderr)
+        return e.returncode, e.stdout if hasattr(e, "stdout") else ""
+
+# ---------- FlashAttention install (best-effort) ----------
+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
+    try:
+        print("Attempting to download and install FlashAttention wheel...")
+        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.")
+    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)
+    if marker.exists():
+        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.")
+        return False
+
+    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
+        marker.write_text("ok")
+        print("download_models.py finished. Marker created.")
+        return True
+    except Exception as e:
+        print("Failed to run download_models.py:", 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
+
+    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.")
+
+    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)
+
+    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)
+        if proc.stderr:
+            print("Inference stderr:", 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)
+        return None
+
+    # locate output video
+    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
+    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).")
+    try:
+        video_path = run_inference(prompt=prompt, image_path=temp_image_path, seed=seed, duration=duration)
+    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).")
+
+        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)
+            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])
+
+    return demo
+
+# ---------- Main entrypoint ----------
+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)
+
+    demo = build_ui()
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))