Update app.py

app.py

@@ -1,48 +1,38 @@
 """
 Sulphur — Image to Video (HF Spaces).
-
-POC-proven pattern:
-- Worker spawned via multiprocessing.get_context("spawn").Process (ZeroGPU
-  patches multiprocessing, so children inherit GPU access)
-- @spaces.GPU(duration=1) writes a signal file; worker inits CUDA within
-  that 1-second window and retains the context for the full generation
-- Triton is stubbed out in the worker so Quanto uses its pure-PyTorch
-  INT8 path — no torch._C._cuda_init() calls after lease expiry
-- Billing: 1 second per generation regardless of inference time
+Clones Wan2GP and downloads models on first run.
+Generation is handled by generate.py called as a subprocess inside @spaces.GPU.
 """
 
-import json
-import multiprocessing
 import os
+import sys
+import subprocess
 import shutil
 import tempfile
 import threading
-import time
+import json
 from pathlib import Path
 
 import gradio as gr
 import spaces
 
-from generate import _worker_entrypoint
-
-_HF_TOKEN     = os.environ.get("HF_TOKEN")
-_PERSISTENT   = Path("/data") if Path("/data").exists() else Path(tempfile.gettempdir())
-WAN2GP_ROOT   = _PERSISTENT / "Wan2GP"
-CKPTS_DIR     = WAN2GP_ROOT / "ckpts"
-LORAS_DIR     = WAN2GP_ROOT / "loras" / "ltx2"
+_HF_TOKEN    = os.environ.get("HF_TOKEN")
+_PERSISTENT  = Path("/data") if Path("/data").exists() else Path(tempfile.gettempdir())
+WAN2GP_ROOT  = _PERSISTENT / "Wan2GP"
+CKPTS_DIR    = WAN2GP_ROOT / "ckpts"
+LORAS_DIR    = WAN2GP_ROOT / "loras" / "ltx2"
 FINETUNES_DIR = WAN2GP_ROOT / "finetunes"
-
-os.environ["WAN2GP_ROOT"] = str(WAN2GP_ROOT)
+GENERATE_PY  = Path(__file__).parent / "generate.py"
 
 SULPHUR_ASSETS = [
     ("SulphurAI/Sulphur-2-base", "sulphur_distil_bf16.safetensors", CKPTS_DIR),
 ]
 LTX_ASSETS = [
     ("SulphurAI/Sulphur-2-base", "experimental/sulphur_experimental_lora_v1.safetensors", LORAS_DIR),
-    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b-distilled-lora-384.safetensors",              LORAS_DIR),
-    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_vae.safetensors",                             CKPTS_DIR),
-    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_text_embedding_projection.safetensors",       CKPTS_DIR),
-    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_embeddings_connector.safetensors",            CKPTS_DIR),
+    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b-distilled-lora-384.safetensors",     LORAS_DIR),
+    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_vae.safetensors",                    CKPTS_DIR),
+    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_text_embedding_projection.safetensors", CKPTS_DIR),
+    ("DeepBeepMeep/LTX-2", "ltx-2.3-22b_embeddings_connector.safetensors",   CKPTS_DIR),
 ]
 
 SULPHUR_FINETUNE = {
@@ -52,14 +42,15 @@ SULPHUR_FINETUNE = {
         "architecture": "ltx2_22B",
         "parent_model_type": "ltx2_22B",
         "description": "LTX-2.3 fine-tuned i2v. Distilled checkpoint.",
+        # Full distilled model — do NOT also preload the rank-768 LoRA (README: use one or the other)
         "URLs": [str(CKPTS_DIR / "sulphur_distil_bf16.safetensors")],
         "preload_URLs": [],
     },
     "num_inference_steps": 8,
     "video_length": 81,
     "resolution": "832x480",
-    "guidance_scale": 5.0,
-    "alt_guidance_scale": 5.0,
+    "guidance_scale": 3.5,
+    "alt_guidance_scale": 3.5,
 }
 
 _setup_lock = threading.Lock()
@@ -69,13 +60,14 @@ _setup_done = False
 def _download(repo_id, filename, dest_dir):
     from huggingface_hub import hf_hub_download
     dest_dir.mkdir(parents=True, exist_ok=True)
-    dest = dest_dir / Path(filename).name
+    dest = dest_dir / Path(filename).name  # flat — strip any subfolder
     if dest.exists():
         print(f"[download] cached: {dest.name}")
         return
     print(f"[download] {repo_id}/{filename}")
     hf_hub_download(repo_id=repo_id, filename=filename,
                     local_dir=str(dest_dir), token=_HF_TOKEN)
+    # hf_hub_download preserves subfolder structure; flatten to dest_dir root
     downloaded = dest_dir / filename
     if downloaded.exists() and not dest.exists():
         shutil.move(str(downloaded), str(dest))
@@ -91,7 +83,6 @@ def setup():
     if not (WAN2GP_ROOT / "shared" / "api.py").exists():
         WAN2GP_ROOT.mkdir(parents=True, exist_ok=True)
         print("[setup] Cloning Wan2GP...")
-        import subprocess
         subprocess.run(
             ["git", "clone", "--depth=1",
              "https://github.com/deepbeepmeep/Wan2GP.git", str(WAN2GP_ROOT)],
@@ -101,9 +92,10 @@ def setup():
     for repo, fname, dest in SULPHUR_ASSETS + LTX_ASSETS:
         _download(repo, fname, dest)
 
+    # Gemma text encoder — must stay in its subfolder (Wan2GP looks there by name)
     _gemma_folder = "gemma-3-12b-it-qat-q4_0-unquantized"
-    _gemma_file   = f"{_gemma_folder}_quanto_bf16_int8.safetensors"
-    gemma_dest    = CKPTS_DIR / _gemma_folder / _gemma_file
+    _gemma_file = f"{_gemma_folder}_quanto_bf16_int8.safetensors"
+    gemma_dest = CKPTS_DIR / _gemma_folder / _gemma_file
     if not gemma_dest.exists():
         from huggingface_hub import hf_hub_download
         print("[download] Gemma text encoder...")
@@ -117,9 +109,7 @@ def setup():
         print("[download] cached: Gemma text encoder")
 
     FINETUNES_DIR.mkdir(parents=True, exist_ok=True)
-    (FINETUNES_DIR / "sulphur_2_base.json").write_text(
-        json.dumps(SULPHUR_FINETUNE, indent=2)
-    )
+    (FINETUNES_DIR / "sulphur_2_base.json").write_text(json.dumps(SULPHUR_FINETUNE, indent=2))
     print("[setup] Done.")
 
 
@@ -128,93 +118,77 @@ setup()
 RESOLUTIONS = ["832x480", "480x832", "640x640", "1024x576", "576x1024"]
 
 
-@spaces.GPU(duration=1)
-def _signal_cuda_init(signal_path):
-    """1-second lease — signals worker to init CUDA, then expires."""
-    Path(signal_path).write_text("go")
-    time.sleep(0.8)
-
-
+@spaces.GPU(duration=80)
 def generate_video(image, prompt, resolution, steps, guidance_scale, frames, seed):
     if image is None:
         raise gr.Error("Please upload an image.")
     if not prompt.strip():
         raise gr.Error("Please enter a prompt.")
 
-    signal_path = tempfile.mktemp(suffix=".signal")
-    result_path = tempfile.mktemp(suffix=".json")
-    out_dir     = tempfile.mkdtemp()
-    out_file    = os.path.join(out_dir, "output.mp4")
-
-    ctx = multiprocessing.get_context("spawn")
-    proc = ctx.Process(
-        target=_worker_entrypoint,
-        args=(image, prompt, out_file, "sulphur-2",
-              int(steps), float(guidance_scale), int(frames),
-              resolution, int(seed), signal_path, result_path),
-        daemon=False,
-    )
-    proc.start()
-
-    log_lines = ["Worker spawned, acquiring GPU lease..."]
-    yield None, "\n".join(log_lines)
-
-    _signal_cuda_init(signal_path)
-
-    log_lines.append("GPU lease expired. Worker retains CUDA context and is generating...")
-    yield None, "\n".join(log_lines)
-
-    last_log_len = 0
-    deadline = time.monotonic() + 600
-
-    while time.monotonic() < deadline:
-        time.sleep(2)
-
-        if os.path.exists(result_path):
-            try:
-                with open(result_path) as f:
-                    data = json.load(f)
-                new_entries = data.get("log", [])[last_log_len:]
-                if new_entries:
-                    log_lines.extend(new_entries)
-                    last_log_len += len(new_entries)
-                yield None, "\n".join(log_lines[-40:])
-                if data.get("done"):
-                    break
-            except Exception:
-                pass
-        else:
-            yield None, "\n".join(log_lines)
-
-    proc.join(timeout=10)
-    if proc.is_alive():
-        proc.terminate()
-
-    for path in (signal_path, result_path):
-        try:
-            os.unlink(path)
-        except Exception:
-            pass
-
-    if os.path.exists(out_file):
-        final = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-        shutil.copy2(out_file, final.name)
-        yield final.name, "\n".join(log_lines) + "\n\n[DONE]"
-    else:
-        yield None, "\n".join(log_lines) + "\n\n[ERROR] No output file produced."
+    out_file = Path(tempfile.mkdtemp()) / "output.mp4"
+    env = {**os.environ, "WAN2GP_ROOT": str(WAN2GP_ROOT)}
+
+    cmd = [
+        sys.executable, str(GENERATE_PY),
+        "--image",  image,
+        "--prompt", prompt,
+        "--output", str(out_file),
+        "--model",  "sulphur-2",
+        "--seed",   str(int(seed)),
+        "--resolution", resolution,
+        "--steps",  str(int(steps)),
+        "--guidance_scale", str(float(guidance_scale)),
+        "--frames", str(int(frames)),
+    ]
+
+    log_lines = []
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                            text=True, bufsize=0, env=env)
+
+    buf = ""
+    while True:
+        chunk = proc.stdout.read(256)
+        if not chunk:
+            break
+        buf += chunk
+        # Split on \r or \n — tqdm uses \r to overwrite progress lines
+        parts = buf.replace("\r", "\n").split("\n")
+        buf = parts[-1]
+        for part in parts[:-1]:
+            stripped = part.strip()
+            if not stripped:
+                continue
+            # Overwrite last line if it looks like a progress bar update
+            if log_lines and ("%" in stripped or "it/s" in stripped or "step" in stripped.lower()):
+                log_lines[-1] = stripped
+            else:
+                log_lines.append(stripped)
+            print(stripped)
+            yield None, "\n".join(log_lines[-30:])
+
+    proc.wait()
+    log = "\n".join(log_lines)
+
+    if proc.returncode != 0 or not out_file.exists():
+        yield None, log + "\n\n[ERROR] Generation failed."
+        return
+
+    final = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    shutil.copy2(out_file, final.name)
+    yield final.name, log + "\n\n[DONE]"
 
 
 with gr.Blocks(title="Sulphur — Image to Video") as demo:
-    gr.Markdown("# Sulphur — Image to Video\nUsing Experimental LoRA v1 + Distillation LoRA")
+    gr.Markdown("# Sulphur — Image to Video\nUsing New Experimental Lora v1")
     with gr.Row():
         with gr.Column(scale=1):
             image_in  = gr.Image(type="filepath", label="Input Image")
             prompt_in = gr.Textbox(label="Prompt", placeholder="Describe the motion…", lines=3)
             with gr.Accordion("Advanced", open=False):
                 resolution_dd = gr.Dropdown(RESOLUTIONS, value="832x480", label="Resolution")
-                steps_sl      = gr.Slider(1, 50, value=8, step=1, label="Steps")
+                steps_sl      = gr.Slider(1, 50, value=8,   step=1,   label="Steps")
                 guidance_sl   = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Guidance Scale")
-                frames_sl     = gr.Slider(17, 257, value=81, step=8, label="Frames")
+                frames_sl     = gr.Slider(17, 257, value=81, step=8,  label="Frames")
                 seed_num      = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
             run_btn = gr.Button("Generate", variant="primary")
         with gr.Column(scale=1):
@@ -228,4 +202,4 @@ with gr.Blocks(title="Sulphur — Image to Video") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+    demo.launch(theme=gr.themes.Soft())