feat: enhance poster generation with offline FLUX support and improve error handling

app.py

@@ -382,13 +382,9 @@ def generate_recap(session_id: str):
     # Generate story
     result = generate_story(packet, session_id=session_id)
 
-    # Try to generate a poster image via BFL FLUX
-    poster_url, poster_status = generate_poster_sync(session_id, result["poster_prompt"])
-    session["poster_url"] = poster_url
-
-    # Format for display
+    # Format for display (poster goes in dedicated section below)
     lines = [
-        "# 📖 Episode Recap\n",
+        "# Episode Recap\n",
         result["short_recap"],
         "",
         "---\n",
@@ -396,11 +392,8 @@ def generate_recap(session_id: str):
         "",
         "---\n",
     ]
-
-    if poster_url:
-        lines.append(f"### 🖼️ Poster\n\n![Recap Poster]({poster_url})\n")
-    else:
-        lines.append(f"### 🎨 Poster Prompt\n```{result['poster_prompt']}```\n\n_{poster_status}_\n")
+    lines.append(f"**Poster prompt:** _{result['poster_prompt']}_\n")
+    lines.append(f"*Click **Generate Poster** below to create an image from this prompt.*\n")
 
     return "\n".join(lines), result
 

app/services/image_gen.py

@@ -1,73 +1,116 @@
 """Black Forest Labs FLUX poster generation module.
 
-Generates a recap poster image from a text prompt using the BFL API.
-When the API key is not set or the call fails, gracefully returns None
-so the pipeline always produces a visible result.
+Generates a recap poster image from a text prompt using FLUX offline
+via ``diffusers`` (same lazy-download pattern as generator.py / minicpm.py).
 
 Sponsor alignment: Black Forest Labs FLUX produces a visible demo artifact
 (the recap poster) rather than being hidden in backend code.
 """
 
 import os
+import uuid
+from pathlib import Path
 from typing import Optional
 
 
 # ── Configuration ───────────────────────────────────────────────────────────
-BFL_API_KEY_ENV = "BFL_API_KEY"
+FLUX_MODEL_ID = "black-forest-labs/FLUX.1-schnell"
+POSTER_DIR = Path("app/assets/posters")
 
 
-def generate_poster(
-    poster_prompt: str,
-    api_key: Optional[str] = None,
-) -> Optional[str]:
-    """Generate a poster image using Black Forest Labs FLUX API.
+# ── Lazy model state ────────────────────────────────────────────────────────
+_model = None
+_model_loaded = False
+_skip_env_var = "CITYQUEST_SKIP_MODEL"
 
-    Args:
-        poster_prompt: Text prompt describing the desired image.
-        api_key: BFL API key. Falls back to ``BFL_API_KEY`` env var.
+
+def _load_model() -> Optional[object]:
+    """Lazy-load the FLUX.1-schnell pipeline via diffusers on first call.
+
+    Uses the Hugging Face Hub cache so subsequent calls are instant.
+    Skips loading if ``CITYQUEST_SKIP_MODEL`` is set (fast-test mode).
 
     Returns:
-        URL string of the generated image, or ``None`` on failure.
+        The FLUX pipeline object, or ``None`` on failure / skip.
     """
-    key = api_key or os.environ.get(BFL_API_KEY_ENV)
-    if not key:
+    global _model, _model_loaded
+    if _model_loaded:
+        return _model
+    _model_loaded = True
+
+    if os.environ.get(_skip_env_var):
+        print("[image_gen] CITYQUEST_SKIP_MODEL set — skipping model load")
+        return None
+
+    try:
+        import torch
+        from diffusers import FluxPipeline
+
+        print(f"[image_gen] Loading {FLUX_MODEL_ID} via diffusers ...")
+        pipe = FluxPipeline.from_pretrained(
+            FLUX_MODEL_ID,
+            torch_dtype=torch.bfloat16,
+        )
+        # Move to GPU if available
+        if torch.cuda.is_available():
+            pipe.to("cuda")
+            print("[image_gen] Pipeline moved to GPU")
+        else:
+            print("[image_gen] Pipeline running on CPU (will be slow)")
+
+        # Enable memory optimisations
+        pipe.enable_model_cpu_offload()
+        _model = pipe
+        print("[image_gen] FLUX pipeline loaded successfully")
+        return pipe
+
+    except ImportError:
         print(
-            "[image_gen] BFL_API_KEY not set — "
-            "set the environment variable to enable poster generation"
+            "[image_gen] Required packages not installed. "
+            "Install with: pip install diffusers transformers accelerate"
         )
         return None
+    except Exception as e:
+        print(f"[image_gen] FLUX model load failed: {type(e).__name__}: {e}")
+        return None
 
-    try:
-        import bfl
 
-        client = bfl.Bfl(api_key=key)
+def generate_poster(poster_prompt: str) -> Optional[str]:
+    """Generate a poster image using offline FLUX.1-schnell.
 
-        # FLUX.1 Pro / FLUX.1.1 Pro — use the latest available
-        result = client.generate_image(
+    Args:
+        poster_prompt: Text prompt describing the desired image.
+
+    Returns:
+        Filesystem path (as string) to the saved image, or ``None`` on failure.
+    """
+    pipe = _load_model()
+    if pipe is None:
+        return None
+
+    try:
+        import torch
+
+        print(f"[image_gen] Generating poster for prompt ({len(poster_prompt)} chars) ...")
+        image = pipe(
             prompt=poster_prompt,
+            guidance_scale=0.0,       # schnell uses 0.0
+            num_inference_steps=4,    # schnell: 4 steps is enough
             width=1024,
             height=768,
-            steps=25,
-            guidance=3.5,
-        )
-
-        # The result object typically has a ``url`` or ``image_url`` attribute
-        image_url = getattr(result, "url", None) or getattr(result, "image_url", None)
-        if image_url:
-            print(f"[image_gen] Poster generated -> {image_url[:80]}...")
-            return image_url
+            generator=torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(42),
+        ).images[0]
 
-        print(f"[image_gen] BFL response had no URL: {result}")
-        return None
+        # Save to disk
+        POSTER_DIR.mkdir(parents=True, exist_ok=True)
+        filename = f"poster_{uuid.uuid4().hex[:8]}.png"
+        filepath = POSTER_DIR / filename
+        image.save(str(filepath))
+        print(f"[image_gen] Poster saved -> {filepath}")
+        return str(filepath)
 
-    except ImportError:
-        print(
-            "[image_gen] 'bfl' Python package not installed. "
-            "Install with: pip install bfl"
-        )
-        return None
     except Exception as e:
-        print(f"[image_gen] BFL API call failed: {type(e).__name__}: {e}")
+        print(f"[image_gen] Poster generation failed: {type(e).__name__}: {e}")
         return None
 
 
@@ -75,33 +118,27 @@ def generate_poster_sync(
     session_id: str,
     poster_prompt: str,
 ) -> tuple[Optional[str], str]:
-    """Generate a poster and return (image_url, status_message).
+    """Generate a poster and return (image_path, status_message).
 
     This is a convenience wrapper for use in the pipeline entry point.
-    The status message explains why no image was generated if the call
-    fails, so the UI can display a helpful fallback.
 
     Args:
         session_id: Session identifier (for logging).
         poster_prompt: Text prompt for the image.
 
     Returns:
-        Tuple of (image_url_or_None, status_message).
+        Tuple of (image_path_or_None, status_message).
     """
-    image_url = generate_poster(poster_prompt)
-
-    if image_url:
-        return image_url, f"\ud83d\uddbc\ufe0f **Poster generated!** [View image]({image_url})"
+    image_path = generate_poster(poster_prompt)
 
-    key = os.environ.get(BFL_API_KEY_ENV)
-    if not key:
-        return None, (
-            "\u26a0\ufe0f Poster generation unavailable — "
-            "set `BFL_API_KEY` environment variable to enable "
-            "Black Forest Labs FLUX poster generation."
+    if image_path:
+        return image_path, (
+            f"**Poster generated!** Saved as `{image_path}` "
+            f"\\u2014 refresh the Recap tab to view."
         )
 
     return None, (
-        "\u26a0\ufe0f Poster generation failed — "
-        "check the BFL API key and try again."
+        "Poster generation unavailable - "
+        "FLUX.1-schnell model could not be loaded. "
+        "Make sure `diffusers` and `transformers` are installed."
     )

app/services/story.py

@@ -312,9 +312,11 @@ def _llm_recap(prompt_template: str, story_packet: dict) -> Optional[dict]:
 def generate_story(story_packet: dict, session_id: Optional[str] = None) -> dict:
     """Generate final recap story from game data and events.
 
-    Attempts Nemotron-based generation first (reuses model already loaded
-    on GPU from game generation), then MiniCPM, then falls back to
-    high-quality templates so the pipeline always returns a result.
+    Uses high-quality template-based generation (deterministic, no GPU
+    required). LLM-based recap paths are disabled by default because:
+      - ``@spaces.GPU`` isolates model cache across subprocess boundaries
+      - Loading a second model instance on the same GPU risks OOM
+      - Templates produce reliable, high-quality output without GPU
 
     Args:
         story_packet: Structured packet with game info, scores, journals, photos.
@@ -324,36 +326,9 @@ def generate_story(story_packet: dict, session_id: Optional[str] = None) -> dict
         Story output dict with keys ``short_recap``, ``long_summary``,
         ``poster_prompt``, ``story_packet``.
     """
-    # Try Nemotron path first (model already loaded on GPU from game generation)
-    llm_result = None
-    template_path = Path("app/prompts/story_recap.txt")
-    prompt_template = ""
-    if template_path.exists():
-        prompt_template = template_path.read_text(encoding="utf-8")
-
-    if prompt_template:
-        llm_result = _llm_recap(prompt_template, story_packet)
-        if llm_result:
-            print("[story] Nemotron recap generated successfully")
-
-    # Fall back to MiniCPM (OpenBMB sponsor model)
-    if not llm_result:
-        try:
-            from app.services.minicpm import generate_recap as minicpm_recap
-            llm_result = minicpm_recap(story_packet)
-            if llm_result:
-                print("[story] MiniCPM recap generated successfully")
-        except Exception as exc:
-            print(f"[story] MiniCPM recap unavailable: {exc}")
-
-    if llm_result and all(k in llm_result for k in ("short_recap", "long_summary", "poster_prompt")):
-        short_recap = llm_result["short_recap"]
-        long_summary = llm_result["long_summary"]
-        poster_prompt = llm_result["poster_prompt"]
-    else:
-        short_recap = _template_short_recap(story_packet)
-        long_summary = _template_long_summary(story_packet)
-        poster_prompt = _template_poster_prompt(story_packet)
+    short_recap = _template_short_recap(story_packet)
+    long_summary = _template_long_summary(story_packet)
+    poster_prompt = _template_poster_prompt(story_packet)
 
     result = {
         "short_recap": short_recap,

requirements.txt

@@ -8,4 +8,5 @@ huggingface_hub
 llama-cpp-python
 gradio>=6.16.0
 jsonschema
+diffusers[torch] transformers accelerate  # uncomment for offline FLUX poster generation
 # modal            # uncomment to deploy to Modal cloud GPU

test_end_to_end.py

@@ -379,17 +379,17 @@ def main():
     print("-" * 80)
 
     try:
-        from app.services.image_gen import generate_poster, generate_poster_sync, BFL_API_KEY_ENV
+        from app.services.image_gen import generate_poster, generate_poster_sync, FLUX_MODEL_ID
 
-        check("Image gen constants loaded", bool(BFL_API_KEY_ENV))
+        check("Image gen constants loaded", bool(FLUX_MODEL_ID))
 
-        # Test with no API key — must return None gracefully
-        result = generate_poster("A test poster prompt", api_key=None)
-        check("generate_poster() with no key returns None", result is None)
+        # Test with skip-model env var — must return None gracefully
+        result = generate_poster("A test poster prompt")
+        check("generate_poster() with skip returns None (no model loaded)", result is None)
 
-        url, status = generate_poster_sync("test-session", "A test poster prompt")
-        check("generate_poster_sync() returns tuple", isinstance(url, type(None)) and isinstance(status, str))
-        check("Status message is informative", "BFL_API_KEY" in status or "Failed" in status or "unavailable" in status)
+        path_or_none, status = generate_poster_sync("test-session", "A test poster prompt")
+        check("generate_poster_sync() returns tuple", isinstance(path_or_none, type(None)) and isinstance(status, str))
+        check("Status message is informative", "unavailable" in status or "Failed" in status)
     except Exception as e:
         check("Image gen module", False, str(e))