feat: Enhance poster generation with FLUX img2img support and improved styling

app.py

@@ -220,6 +220,7 @@ def use_hint(session_id: str, task_id: str, team_id: str = "team-a"):
     return f"💡 Hint used for {task_id} (−5 pts)"
 
 
+@spaces.GPU(duration=120)
 def record_journal(
     session_id: str,
     transcript: str = "",
@@ -231,6 +232,11 @@ def record_journal(
 ):
     """Record a journal entry, summarize it, and return the result.
 
+    Decorated with ``@spaces.GPU`` so the Cohere ASR model can run on the
+    GPU. On HF ZeroGPU, ``torch.cuda.is_available()`` is ``False`` outside
+    a ``@spaces.GPU`` function, so without this the ASR model would load
+    on CPU and the voice-journal button would appear to hang.
+
     Two input paths are supported:
 
     * **Voice path** — pass ``audio_path`` (e.g. a path returned by

app/services/image_gen.py

@@ -24,9 +24,20 @@ POSTER_DIR = Path("app/assets/posters")
 # ── Lazy model state ────────────────────────────────────────────────────────
 _model = None
 _model_loaded = False
+_img2img_model = None
 _skip_env_var = "CITYQUEST_SKIP_MODEL"
 
 
+# ── Poster styling ───────────────────────────────────────────────────────────
+# Appended to every FLUX prompt so the output reads as a polished poster
+# rather than a literal snapshot.
+POSTER_STYLE = (
+    " Illustrated travel-poster art style, bold dynamic composition, "
+    "rich saturated colors, dramatic lighting, clean graphic shapes, "
+    "high detail, trending on ArtStation. No text, no watermark."
+)
+
+
 def _load_model() -> Optional[object]:
     """Lazy-load the FLUX.1-schnell pipeline via diffusers on first call.
 
@@ -137,7 +148,7 @@ def generate_poster(poster_prompt: str) -> Optional[str]:
 
         print(f"[image_gen] Generating poster for prompt ({len(poster_prompt)} chars) ...")
         image = pipe(
-            prompt=poster_prompt,
+            prompt=poster_prompt + POSTER_STYLE,
             guidance_scale=0.0,       # schnell uses 0.0
             num_inference_steps=4,    # schnell: 4 steps is enough
             width=1024,
@@ -145,19 +156,101 @@ def generate_poster(poster_prompt: str) -> Optional[str]:
             generator=torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(42),
         ).images[0]
 
-        # 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)
+        return _save_poster(image)
 
     except Exception as e:
         print(f"[image_gen] Poster generation failed: {type(e).__name__}: {e}")
         return None
 
 
+def _load_img2img_model() -> Optional[object]:
+    """Lazy-load a FLUX image-to-image pipeline, reusing the text-to-image
+    model's weights so we don't load a second ~24 GB copy.
+
+    Returns:
+        A ``FluxImg2ImgPipeline``, or ``None`` on failure / skip.
+    """
+    global _img2img_model
+    if _img2img_model is not None:
+        return _img2img_model
+
+    pipe = _load_model()  # the base text-to-image FLUX pipeline
+    if pipe is None:
+        return None
+
+    try:
+        from diffusers import FluxImg2ImgPipeline
+
+        # from_pipe reuses the already-loaded components (and any CPU-offload
+        # hooks) instead of downloading / instantiating a second model.
+        _img2img_model = FluxImg2ImgPipeline.from_pipe(pipe)
+        print("[image_gen] FLUX img2img pipeline ready")
+        return _img2img_model
+    except Exception as e:
+        print(f"[image_gen] FLUX img2img init failed: {type(e).__name__}: {e}")
+        return None
+
+
+def generate_poster_img2img(
+    poster_prompt: str,
+    photo_path: str,
+    strength: float = 0.8,
+) -> Optional[str]:
+    """Generate a poster by letting FLUX re-imagine an uploaded gameplay photo.
+
+    The photo seeds the composition/colors; the prompt drives the poster
+    look. ``strength`` controls how far FLUX moves from the photo
+    (0 = identical, 1 = ignore photo). ~0.8 keeps the scene recognizable
+    while turning it into a stylized poster.
+
+    Args:
+        poster_prompt: Text prompt describing the desired poster.
+        photo_path: Path to an uploaded gameplay photo to use as the seed.
+        strength: Denoising strength in ``[0, 1]``.
+
+    Returns:
+        Filesystem path to the saved poster, or ``None`` on failure.
+    """
+    if not Path(photo_path).exists():
+        print(f"[image_gen] img2img photo not found: {photo_path}")
+        return None
+
+    pipe = _load_img2img_model()
+    if pipe is None:
+        return None
+
+    try:
+        import torch
+        from PIL import Image
+
+        init_image = Image.open(photo_path).convert("RGB").resize((1024, 768))
+        print(f"[image_gen] Re-imagining {photo_path} via FLUX img2img (strength={strength}) ...")
+        image = pipe(
+            prompt=poster_prompt + POSTER_STYLE,
+            image=init_image,
+            strength=strength,
+            guidance_scale=0.0,        # schnell uses 0.0
+            num_inference_steps=8,     # effective steps ≈ steps * strength
+            generator=torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(42),
+        ).images[0]
+
+        return _save_poster(image)
+
+    except Exception as e:
+        print(f"[image_gen] img2img poster generation failed: {type(e).__name__}: {e}")
+        return None
+
+
+def _save_poster(image) -> str:
+    """Save a PIL image to the poster directory and return its path."""
+    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)
+
+
 # ── Photo collage poster (uses uploaded gameplay photos) ─────────────────
 
 def generate_collage_poster(
@@ -308,9 +401,13 @@ def generate_poster_sync(
 ) -> tuple[Optional[str], str]:
     """Generate a poster and return (image_path, status_message).
 
-    Strategy:
-      1. If uploaded photos exist → build a photo collage (instant, no GPU).
-      2. Otherwise → fall back to FLUX text-to-image (needs GPU).
+    Strategy (in priority order):
+      1. If a gameplay photo exists → **FLUX image-to-image**: the model
+         re-imagines the photo into a stylized poster (GPU). This is the
+         desired behavior — a *generated* poster seeded by a real moment.
+      2. No photo → **FLUX text-to-image** from the prompt (GPU).
+      3. FLUX unavailable (offline / no GPU) → Pillow photo collage as a
+         graceful fallback so the demo still produces something.
 
     Args:
         session_id: Session identifier (for logging).
@@ -320,26 +417,39 @@ def generate_poster_sync(
     Returns:
         Tuple of (image_path_or_None, status_message).
     """
-    # ── Strategy 1: Photo collage from uploaded gameplay photos ─────────
-    if photo_paths:
-        print(f"[image_gen] Attempting collage from {len(photo_paths)} photo(s)")
-        collage = generate_collage_poster(photo_paths)
-        if collage:
-            return collage, (
-                f"**Poster created from your gameplay photos!** 📸\n"
-                f"Saved as `{collage}` — refresh the Recap tab to view."
+    valid_photos = [p for p in (photo_paths or []) if p and Path(p).exists()]
+
+    # ── Strategy 1: FLUX img2img from an uploaded gameplay photo ─────────
+    if valid_photos:
+        print(f"[image_gen] Attempting FLUX img2img from {len(valid_photos)} photo(s)")
+        image_path = generate_poster_img2img(poster_prompt, valid_photos[0])
+        if image_path:
+            return image_path, (
+                f"**AI poster generated from your photo!** 🎨\n"
+                f"FLUX re-imagined your gameplay shot into a poster — "
+                f"saved as `{image_path}`."
             )
-        print("[image_gen] Collage failed, falling back to FLUX text-to-image")
+        print("[image_gen] img2img failed, falling back to FLUX text-to-image")
 
-    # ── Strategy 2: FLUX text-to-image (GPU required) ─────────────────
+    # ── Strategy 2: FLUX text-to-image from the prompt ─────────────────
     image_path = generate_poster(poster_prompt)
-
     if image_path:
         return image_path, (
             f"**Poster generated!** Saved as `{image_path}` "
-            f"\\u2014 refresh the Recap tab to view."
+            f"— refresh the Recap tab to view."
         )
 
+    # ── Strategy 3: Pillow collage fallback (offline / no GPU) ─────────
+    if valid_photos:
+        print("[image_gen] FLUX unavailable, falling back to photo collage")
+        collage = generate_collage_poster(valid_photos)
+        if collage:
+            return collage, (
+                f"**Poster created from your gameplay photos.** 📸\n"
+                f"_(FLUX was unavailable, so this is a collage of your photos "
+                f"rather than an AI-generated poster.)_ Saved as `{collage}`."
+            )
+
     return None, (
         "⚠️ **Poster generation unavailable** — FLUX.1-schnell model could not be loaded.\n\n"
         "Possible causes:\n"