Update app.py

app.py

@@ -1,4 +1,5 @@
-# app.py ──────────────────────────────────────────────────────────────────────
+# app.py — Robust Character Generator Space
+
 from datasets import load_dataset
 import gradio as gr, json, os, random, torch, spaces
 from diffusers import FluxPipeline, AutoencoderKL
@@ -7,10 +8,10 @@ from live_preview_helpers import (
     flux_pipe_call_that_returns_an_iterable_of_images as flux_iter,
 )
 
-# ─────────────────────────── 1.  Device ─────────────────────────────────────
+# 1. Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ────────────────────── 2.  FLUX image pipeline ─────────────────────────────
+# 2. FLUX image pipeline
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
 ).to(device)
@@ -19,7 +20,7 @@ good_vae = AutoencoderKL.from_pretrained(
 ).to(device)
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_iter.__get__(pipe)
 
-# ───────────────────────── 3.  LLM client (robust) ──────────────────────────
+# 3. LLM client (robust)
 LLM_SPACES = [
     "https://huggingfaceh4-zephyr-chat.hf.space",
     "meta-llama/Llama-3.3-70B-Instruct",
@@ -27,14 +28,11 @@ LLM_SPACES = [
 ]
 
 def first_live_space(space_ids: list[str]) -> Client:
-    """
-    Return the first Space whose /chat endpoint answers a 1-token echo.
-    """
     for sid in space_ids:
         try:
             print(f"[info] probing {sid}")
             c = Client(sid, hf_token=os.getenv("HF_TOKEN"))
-            _ = c.predict("ping", 8, api_name="/chat")  # simple health check
+            _ = c.predict("ping", 8, api_name="/chat")
             print(f"[info] using {sid}")
             return c
         except Exception as e:
@@ -42,18 +40,12 @@ def first_live_space(space_ids: list[str]) -> Client:
     raise RuntimeError("No live chat Space found!")
 
 llm_client = first_live_space(LLM_SPACES)
-CHAT_API   = "/chat"  # universal endpoint for TGI-style Spaces
+CHAT_API = "/chat"
 
 def call_llm(prompt: str,
              max_tokens: int = 256,
              temperature: float = 0.6,
              top_p: float = 0.9) -> str:
-    """
-    Send a single-message chat to the Space. Extra sliders in the remote UI must
-    be supplied in positional order after the prompt, so we match Zephyr/Gemma:
-        [prompt, max_tokens, temperature, top_p, repeat_penalty, presence_penalty]
-    We pass only the first four; the Space will fill the rest with defaults.
-    """
     try:
         return llm_client.predict(
             prompt, max_tokens, temperature, top_p, api_name=CHAT_API
@@ -62,26 +54,21 @@ def call_llm(prompt: str,
         print(f"[error] LLM failure → {exc}")
         return "…"
 
-# ──────────────────────── 4.  Persona dataset ──────────────────────────────
+# 4. Persona dataset
 ds = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")
 
 def random_persona() -> str:
     return ds[random.randint(0, len(ds) - 1)]["persona"]
 
-# ─────────────────────────── 5.  Text prompts ───────────────────────────────
+# 5. Text prompts
 PROMPT_TEMPLATE = """Generate a character with this persona description:
-
 {persona_description}
-
 In a world with this description:
-
 {world_description}
-
 Write the character in JSON with keys:
 name, background, appearance, personality, skills_and_abilities,
 goals, conflicts, backstory, current_situation,
 spoken_lines (list of strings).
-
 Respond with JSON only (no markdown)."""
 
 WORLD_PROMPT = (
@@ -89,12 +76,25 @@ WORLD_PROMPT = (
     "Respond with the description only."
 )
 
-# ───────────────────────── 6.  Helper functions ─────────────────────────────
+# 6. Helper functions
 def random_world() -> str:
     return call_llm(WORLD_PROMPT, max_tokens=120)
 
+def safe_json_parse(raw):
+    """Try to parse JSON, return None if fail, and log."""
+    try:
+        return json.loads(raw)
+    except Exception as e:
+        print(f"[ERROR] JSON parsing failed: {e}")
+        print(f"[DEBUG] Raw output: {raw[:1000]}")
+        return None
+
 @spaces.GPU(duration=75)
 def infer_flux(character_json):
+    # Defensive: If not a dict or missing appearance, bail out
+    if not isinstance(character_json, dict) or "appearance" not in character_json:
+        print("[ERROR] No valid appearance to generate image.")
+        return None
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=character_json["appearance"],
         guidance_scale=3.5,
@@ -109,6 +109,7 @@ def infer_flux(character_json):
 
 def generate_character(world_desc: str, persona_desc: str,
                        progress=gr.Progress(track_tqdm=True)):
+    # First attempt
     raw = call_llm(
         PROMPT_TEMPLATE.format(
             persona_description=persona_desc,
@@ -116,25 +117,35 @@ def generate_character(world_desc: str, persona_desc: str,
         ),
         max_tokens=1024,
     )
-    try:
-        return json.loads(raw)
-    except json.JSONDecodeError:
-        # retry once if the model didn’t return valid JSON
-        raw = call_llm(
-            PROMPT_TEMPLATE.format(
-                persona_description=persona_desc,
-                world_description=world_desc,
-            ),
-            max_tokens=1024,
-        )
-        return json.loads(raw)
+    character = safe_json_parse(raw)
+    if character:
+        return character
 
-# ─────────────────────────── 7.  Gradio UI ──────────────────────────────────
+    # Retry once
+    raw2 = call_llm(
+        PROMPT_TEMPLATE.format(
+            persona_description=persona_desc,
+            world_description=world_desc,
+        ),
+        max_tokens=1024,
+    )
+    character2 = safe_json_parse(raw2)
+    if character2:
+        return character2
+
+    # If both fail, return error and raw outputs for debugging
+    return {
+        "error": "LLM did not return valid JSON after 2 attempts.",
+        "first_raw": raw,
+        "second_raw": raw2,
+        "tip": "Check your LLM prompt and output. Try regenerating.",
+    }
+
+# 7. Gradio UI
 DESCRIPTION = """
 * Generates a JSON character sheet from a world + persona.
 * Appearance images via **FLUX-dev**; story text via Zephyr-chat or Gemma fallback.
 * Personas sampled from **FinePersonas-Lite**.
-
 Tip → Shuffle the world then persona for rapid inspiration.
 """