Run on ZeroGPU: spaces.GPU session generator + command queue for controls

README.md

@@ -15,4 +15,5 @@ Interactive [ABot-World](https://huggingface.co/YiYiXu/ABot-World-0-5B-LF-Diffus
 action-conditioned world model, Wan2.2-TI2V-5B) running natively on Modular Diffusers: the whole app is a
 `pipe.stream(action_source=...)` consumer — the pipeline polls the held control once per generated block.
 
-Needs a persistent GPU (the rollout keeps a K/V cache alive across blocks); ~20 GB VRAM in bf16.
+Runs on ZeroGPU: one `@spaces.GPU` generator owns the whole session (the rollout keeps a K/V cache alive
+across blocks) and the held control reaches it through a queue. ~20 GB VRAM in bf16.

app.py

@@ -1,3 +1,7 @@
+import spaces  # must come before torch: on ZeroGPU it patches CUDA so the model can be staged at import
+
+import multiprocessing
+import queue
 import threading
 
 import gradio as gr
@@ -12,7 +16,7 @@ REPO = "YiYiXu/ABot-World-0-5B-LF-Diffusers"
 
 pipe = ABotWorldStreamingBlocks().init_pipeline(REPO)
 pipe.load_components(dtype=torch.bfloat16)
-pipe.to("cuda")
+pipe.to("cuda")  # on ZeroGPU this only packs the weights; they land on the GPU inside the @spaces.GPU call
 
 DEFAULT_PROMPT = (
     "A realistic outdoor world scene with a navigable path, natural lighting, "
@@ -31,45 +35,63 @@ CONTROLS = {
     "· idle": [0] * 8,
 }
 
-# one live world at a time: buttons write the held control here, the rollout polls it per block
-session = {"action": [0] * 8, "stop": False}
+# One live world at a time. On ZeroGPU the rollout runs in a forked worker process, so the held control
+# crosses over a multiprocessing queue created before the fork: buttons put an action, Stop puts None,
+# and the rollout drains the queue once per block.
+commands = multiprocessing.Queue()
 run_lock = threading.Lock()
 
 
 def hold(control):
-    session["action"] = CONTROLS[control]
+    commands.put(CONTROLS[control])
     return f"holding: {control}"
 
 
 def stop():
-    session["stop"] = True
+    commands.put(None)
     return "stopping after this block…"
 
 
+@spaces.GPU(duration=lambda image, prompt, max_blocks: 30 + 4 * int(max_blocks))
+def rollout(image, prompt, max_blocks):
+    held = CONTROLS["▲ forward (W)"]
+
+    def action_source(block_index):
+        nonlocal held
+        while True:
+            try:
+                command = commands.get_nowait()
+            except queue.Empty:
+                break
+            if command is None:
+                return None
+            held = command
+        return held if block_index < max_blocks else None
+
+    for event in pipe.stream(
+        prompt=(prompt or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT,
+        image=image,
+        action_source=action_source,
+        generator=torch.Generator("cpu").manual_seed(42),
+    ):
+        if event.path != "denoise.rollout":
+            continue
+        frames = event.state.get("frames")
+        yield (frames[-1] * 255).clip(0, 255).astype(np.uint8)
+
+
 def play(image, prompt, max_blocks):
     if not run_lock.acquire(blocking=False):
         raise gr.Error("A session is already running — stop it first.")
     try:
+        while True:  # drop controls left over from the previous session
+            try:
+                commands.get_nowait()
+            except queue.Empty:
+                break
         if image is None:
             image = load_image("examples/desert_valley.png")
-        session["stop"] = False
-        session["action"] = CONTROLS["▲ forward (W)"]
-
-        def action_source(block_index):
-            if session["stop"] or block_index >= max_blocks:
-                return None
-            return session["action"]
-
-        for event in pipe.stream(
-            prompt=(prompt or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT,
-            image=image,
-            action_source=action_source,
-            generator=torch.Generator("cpu").manual_seed(42),
-        ):
-            if event.path != "denoise.rollout":
-                continue
-            frames = event.state.get("frames")
-            yield (frames[-1] * 255).clip(0, 255).astype(np.uint8)
+        yield from rollout(image, prompt, max_blocks)
     finally:
         run_lock.release()
 
@@ -86,7 +108,7 @@ per block ([huggingface/diffusers#14159](https://github.com/huggingface/diffuser
         with gr.Column(scale=1):
             image = gr.Image(label="starting scene", type="pil", height=220)
             prompt = gr.Textbox(label="scene prompt", value=DEFAULT_PROMPT, lines=2)
-            max_blocks = gr.Slider(label="max blocks (~seconds)", minimum=4, maximum=120, step=1, value=60)
+            max_blocks = gr.Slider(label="max blocks (~seconds)", minimum=4, maximum=40, step=1, value=20)
             with gr.Row():
                 start_button = gr.Button("Start", variant="primary")
                 stop_button = gr.Button("Stop")