Upload app.py

app.py

@@ -99,141 +99,153 @@ def warmup_models():
 
 # ── Frame 1: text ─────────────────────────────────────────────────────────────
 
-@spaces.GPU(duration=80)
-def run_text_gen(narrator: str):
-    """GPU frame 1: GGUF text model → room JSON + narrations.
-    Returns a dict stored in gen_state; no PIL images yet.
+@spaces.GPU(duration=360)  # text + images + TTS combined
+def run_generation_pipeline(narrator: str):
     """
-    # Normalise: strip whitespace and fall back to the first listed voice
+    Single GPU frame:
+        GGUF -> room generation
+        FLUX -> images
+        VoxCPM -> audio
+    """
+
     narrator = (narrator or "").strip()
     if narrator not in misc_f.VOICES:
         narrator = next(iter(misc_f.VOICES))
-        log.warning("run_text_gen: invalid narrator, using %r", narrator)
-    log.info("run_text_gen: narrator = %r", narrator)
-    log.info("run_text_gen: loading GGUF text model…")
+
+    # ------------------------------------------------------------------
+    # TEXT
+    # ------------------------------------------------------------------
+
+    log.info("Loading text model...")
     llm = misc_f.load_text_model()
 
-    room  = misc_f.generate_game(llm)
+    room = misc_f.generate_game(llm)
     chain = generate_fixed_chain()
     narrations = misc_f.pregenerate_text(llm, room, chain)
 
     misc_f.unload_text_model(llm)
-    log.info("run_text_gen: done.")
-    return {"room": room, "chain": chain, "narrator": narrator,
-            "narrations": narrations}
 
+    # ------------------------------------------------------------------
+    # IMAGES
+    # ------------------------------------------------------------------
 
-def on_text_done(gen_state: dict):
-    """CPU step: text is ready — do NOT show main screen yet.
-    Keep the generating overlay up; images and audio still need to finish.
-    We stash the story in the narrative box (hidden) so on_all_done can
-    populate it in one shot when everything is ready.
-    """
-    story = gen_state["room"].get("room_story", "")
-    log.info("on_text_done: text ready, waiting for images + audio before showing screen.")
-    # Return story so it's available to on_all_done via narrative state,
-    # but leave both screens unchanged.
-    return story   # → narrative (hidden on main_screen, harmless update)
+    log.info("Loading image model...")
+    image_model = misc_f.load_image_model()
 
+    room_img = misc_f.generate_image(
+        image_model,
+        room["room_prompt"],
+        "room"
+    )
 
-# ── Frame 2: images + TTS (merged to avoid ZeroGPU rate-limit on 3rd lease) ───
-#
-# Running FLUX then VoxCPM sequentially in one @spaces.GPU frame means we only
-# consume 2 GPU leases total (text + this one) instead of 3, which keeps us
-# under the ZeroGPU 429 rate-limit.  Each model is fully unloaded before the
-# next one is loaded so peak VRAM stays the same as before.
-
-@spaces.GPU(duration=360)
-def run_images_and_tts(gen_state: dict):
-    """GPU frame 2: FLUX → 9 images, then VoxCPM → audio clips.
-    Returns a combined dict with both 'img' and 'audio' keys.
-    """
-    room       = gen_state["room"]
-    chain      = gen_state["chain"]
-    narrations = gen_state["narrations"]
-    narrator   = gen_state["narrator"]
+    loc_imgs = [
+        misc_f.generate_image(
+            image_model,
+            room["containers"][i]["container_prompt"],
+            "location"
+        )
+        for i in range(4)
+    ]
 
-    # ── FLUX ──────────────────────────────────────────────────────────────────
-    log.info("run_images_and_tts: loading FLUX image model…")
-    image_model = misc_f.load_image_model()
+    door_img = misc_f.generate_image(
+        image_model,
+        room["door_prompt"],
+        "location"
+    )
 
-    room_img = misc_f.generate_image(image_model, room["room_prompt"],                               "room")
-    loc0     = misc_f.generate_image(image_model, room["containers"][0]["container_prompt"],         "location")
-    loc1     = misc_f.generate_image(image_model, room["containers"][1]["container_prompt"],         "location")
-    loc2     = misc_f.generate_image(image_model, room["containers"][2]["container_prompt"],         "location")
-    loc3     = misc_f.generate_image(image_model, room["containers"][3]["container_prompt"],         "location")
-    door_img = misc_f.generate_image(image_model, room["door_prompt"],                               "location")
-    key0     = misc_f.generate_image(image_model, room["keys"][0]["key_prompt"],                     "item")
-    key1     = misc_f.generate_image(image_model, room["keys"][1]["key_prompt"],                     "item")
-    door_key = misc_f.generate_image(image_model, room["door_key_prompt"],                           "item")
+    key0 = misc_f.generate_image(
+        image_model,
+        room["keys"][0]["key_prompt"],
+        "item"
+    )
+
+    key1 = misc_f.generate_image(
+        image_model,
+        room["keys"][1]["key_prompt"],
+        "item"
+    )
+
+    door_key = misc_f.generate_image(
+        image_model,
+        room["door_key_prompt"],
+        "item"
+    )
 
     misc_f.unload_image_model(image_model)
-    log.info("run_images_and_tts: FLUX done, unloaded.")
 
-    img_result = {
-        "room_img": room_img,
-        "loc_imgs": [loc0, loc1, loc2, loc3],
-        "door_img": door_img,
-        "key_imgs": [key0, key1, door_key],
-    }
+    # ------------------------------------------------------------------
+    # TTS
+    # ------------------------------------------------------------------
+
+    log.info("Loading TTS model...")
+    tts_model = misc_f.load_tts_model()
+
+    audio_paths = misc_f.pregenerate_audio(
+        tts_model,
+        room,
+        chain,
+        narrations,
+        narrator,
+    )
 
-    # ── VoxCPM ────────────────────────────────────────────────────────────────
-    log.info("run_images_and_tts: loading VoxCPM TTS model…")
-    tts_model   = misc_f.load_tts_model()
-    audio_paths = misc_f.pregenerate_audio(tts_model, room, chain, narrations, narrator)
     misc_f.unload_tts_model(tts_model)
-    log.info("run_images_and_tts: VoxCPM done, unloaded.")
 
-    return {"img": img_result, "audio": audio_paths}
+    # ------------------------------------------------------------------
+    # RETURN EVERYTHING
+    # ------------------------------------------------------------------
+
+    return {
+        "room": room,
+        "chain": chain,
+        "narrator": narrator,
+        "narrations": narrations,
+        "images": {
+            "room_img": room_img,
+            "loc_imgs": loc_imgs,
+            "door_img": door_img,
+            "key_imgs": [key0, key1, door_key],
+        },
+        "audio_paths": audio_paths,
+    }
 
 
-def on_all_done(gen_state: dict, combined: dict):
-    """CPU step: everything is ready — flip to main screen and populate all widgets.
-    'combined' is the merged dict returned by run_images_and_tts.
-    """
-    room        = gen_state["room"]
-    chain       = gen_state["chain"]
-    narrator    = gen_state["narrator"]
-    narrations  = gen_state["narrations"]
-    img_result  = combined["img"]
-    audio_paths = combined["audio"]
-    loc_imgs    = img_result["loc_imgs"]
-    key_imgs    = img_result["key_imgs"]
+def on_generation_done(result):
+    room = result["room"]
 
     state = {
-        "room":         room,
-        "chain":        chain,
-        "loc_imgs":     loc_imgs,
-        "key_imgs":     key_imgs,
-        "found_keys":   set(),
+        "room": room,
+        "chain": result["chain"],
+        "loc_imgs": result["images"]["loc_imgs"],
+        "key_imgs": result["images"]["key_imgs"],
+        "found_keys": set(),
         "tried_combos": set(),
-        "door_opened":  False,
-        "narrator":     narrator,
-        "narrations":   narrations,
-        "audio_paths":  audio_paths,
+        "door_opened": False,
+        "narrator": result["narrator"],
+        "narrations": result["narrations"],
+        "audio_paths": result["audio_paths"],
     }
 
-    story           = room.get("room_story", "")
-    room_story_path = audio_paths.get("room_story")
-    log.info("on_all_done: everything ready — switching to main screen.")
+    story = room.get("room_story", "")
+
     return (
-        gr.update(visible=False),   # gen_screen
-        gr.update(visible=True),    # main_screen
-        story,                      # narrative
-        img_result["room_img"],     # main_img
-        loc_imgs,                   # gallery_top
-        img_result["door_img"],     # mid_img
-        [],                         # stack_gallery (empty inventory at start)
-        None,                       # small_img
-        state,                      # room_state
-        -1,                         # sel_container
-        "",                         # sel_key
-        gr.update(visible=True),    # door_btn
-        "Container: \u2014\nKey: \u2014",  # selection_display
-        gr.update(value=room_story_path, visible=bool(room_story_path)),  # narration_audio
+        gr.update(visible=False),
+        gr.update(visible=True),
+        story,
+        result["images"]["room_img"],
+        result["images"]["loc_imgs"],
+        result["images"]["door_img"],
+        [],
+        None,
+        state,
+        -1,
+        "",
+        gr.update(visible=True),
+        "Container: —\nKey: —",
+        gr.update(
+            value=result["audio_paths"].get("room_story"),
+            visible=bool(result["audio_paths"].get("room_story"))
+        ),
     )
-
-
 # ─── Gallery selection handlers ───────────────────────────────────────────���───
 
 def on_container_select_full(evt: gr.SelectData, state, key_id):
@@ -592,41 +604,39 @@ with gr.Blocks(title="1000 Rooms") as demo:
     # on_all_done hides the overlay via JS once everything is ready.
 
     gen_btn.click(
-        fn=run_text_gen,
-        inputs=[narrator_state],
-        outputs=[gen_state],
-        js="""() => {
-            const ov = document.getElementById('nn-overlay');
-            if (ov) ov.classList.remove('hidden');
-            startNN();
-        }""",
-    ).then(
-        fn=on_text_done,
-        inputs=[gen_state],
-        outputs=[narrative],   # stash story; overlay stays up
-    ).then(
-        fn=run_images_and_tts,
-        inputs=[gen_state],
-        outputs=[combined_state],
-    ).then(
-        fn=on_all_done,
-        inputs=[gen_state, combined_state],
-        outputs=[
-            gen_screen, main_screen,
-            narrative,
-            main_img, gallery_top, mid_img,
-            stack_gallery, small_img,
-            room_state,
-            sel_container, sel_key,
-            door_btn, selection_display,
-            narration_audio,
-        ],
-        js="""() => {
-            const ov = document.getElementById('nn-overlay');
-            if (ov) ov.classList.add('hidden');
-            stopNN();
-        }""",
-    )
+    fn=run_generation_pipeline,
+    inputs=[narrator_state],
+    outputs=[combined_state],
+    js="""() => {
+        const ov = document.getElementById('nn-overlay');
+        if (ov) ov.classList.remove('hidden');
+        startNN();
+    }"""
+).then(
+    fn=on_generation_done,
+    inputs=[combined_state],
+    outputs=[
+        gen_screen,
+        main_screen,
+        narrative,
+        main_img,
+        gallery_top,
+        mid_img,
+        stack_gallery,
+        small_img,
+        room_state,
+        sel_container,
+        sel_key,
+        door_btn,
+        selection_display,
+        narration_audio,
+    ],
+    js="""() => {
+        const ov = document.getElementById('nn-overlay');
+        if (ov) ov.classList.add('hidden');
+        stopNN();
+    }"""
+)
 
 
     # ── Gallery selection wiring ──────────────────────────────────────────────