Update app.py

app.py

@@ -1,87 +1,74 @@
+# app.py ──────────────────────────────────────────────────────────────────────
 from datasets import load_dataset
-import gradio as gr
-from gradio_client import Client
-import json, os, random, torch, spaces
+import gradio as gr, json, os, random, torch, spaces
 from diffusers import FluxPipeline, AutoencoderKL
-from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+from gradio_client import Client
+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.  Image / FLUX pipeline ───────────────────────────
+# ────────────────────── 2.  FLUX image pipeline ─────────────────────────────
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
 ).to(device)
 good_vae = AutoencoderKL.from_pretrained(
     "black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16
 ).to(device)
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = (
-    flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
-)
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_iter.__get__(pipe)
 
-# ───────────────────────── 3.  LLM client (robust) ───────────────────────────
-def _first_working_client(candidates: list[str]) -> Client:
+# ───────────────────────── 3.  LLM client (robust) ──────────────────────────
+LLM_SPACES = [
+    "https://huggingfaceh4-zephyr-chat.hf.space",
+    "HuggingFaceH4/zephyr-chat",
+    "huggingface-projects/gemma-2-9b-it",
+]
+
+def first_live_space(space_ids: list[str]) -> Client:
     """
-    Try a list of Space URLs / repo-ids, return the first that gives a JSON config.
+    Return the first Space whose /chat endpoint answers a 1-token echo.
     """
-    for src in candidates:
+    for sid in space_ids:
         try:
-            print(f"[info] Trying LLM Space: {src}")
-            c = Client(src, hf_token=os.getenv("HF_TOKEN"))  # token optional
-            # If this passes, the config was parsed as JSON
-            c.view_api()
-            print(f"[info] Selected LLM Space: {src}")
+            print(f"[info] probing {sid}")
+            c = Client(sid, hf_token=os.getenv("HF_TOKEN"))
+            _ = c.predict("ping", 8, api_name="/chat")  # simple health check
+            print(f"[info] using {sid}")
             return c
         except Exception as e:
-            print(f"[warn] {src} not usable → {e}")
-    raise RuntimeError("No usable LLM Space found!")
+            print(f"[warn] {sid} unusable → {e}")
+    raise RuntimeError("No live chat Space found!")
 
-LLM_CANDIDATES = [
-    "https://huggingfaceh4-zephyr-chat.hf.space",  # direct URL
-    "HuggingFaceH4/zephyr-chat",                  # repo slug
-    "huggingface-projects/gemma-2-9b-it",         # fallback Space
-]
+llm_client = first_live_space(LLM_SPACES)
+CHAT_API   = "/chat"  # universal endpoint for TGI-style Spaces
 
-llm_client = _first_working_client(LLM_CANDIDATES)
-CHAT_API   = llm_client.view_api()[0]["api_name"]  # safest way to get endpoint
-
-def call_llm(
-    user_prompt: str,
-    system_prompt: str = "You are a helpful creative assistant.",
-    history: list | None = None,
-    temperature: float = 0.7,
-    top_p: float = 0.9,
-    max_tokens: int = 1024,
-) -> str:
+def call_llm(prompt: str,
+             max_tokens: int = 256,
+             temperature: float = 0.6,
+             top_p: float = 0.9) -> str:
     """
-    Unified chat wrapper – works for both Zephyr and Gemma Spaces.
+    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.
     """
-    history = history or []
     try:
-        result = llm_client.predict(
-            user_prompt,
-            system_prompt,
-            history,
-            temperature,
-            top_p,
-            max_tokens,
-            api_name=CHAT_API,
-        )
-        # Some Spaces return string, some return (…, history) tuple
-        if isinstance(result, str):
-            return result.strip()
-        return result[1][0][-1].strip()
-    except Exception as e:
-        print(f"[error] LLM call failed → {e}")
+        return llm_client.predict(
+            prompt, max_tokens, temperature, top_p, api_name=CHAT_API
+        ).strip()
+    except Exception as exc:
+        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.  Prompt templates ───────────────────────────
+# ─────────────────────────── 5.  Text prompts ───────────────────────────────
 PROMPT_TEMPLATE = """Generate a character with this persona description:
 
 {persona_description}
@@ -102,9 +89,9 @@ WORLD_PROMPT = (
     "Respond with the description only."
 )
 
-# ─────────────────────── 6.  Helper functions ───────────────────────────────
+# ───────────────────────── 6.  Helper functions ─────────────────────────────
 def random_world() -> str:
-    return call_llm(WORLD_PROMPT)
+    return call_llm(WORLD_PROMPT, max_tokens=120)
 
 @spaces.GPU(duration=75)
 def infer_flux(character_json):
@@ -132,7 +119,7 @@ def generate_character(world_desc: str, persona_desc: str,
     try:
         return json.loads(raw)
     except json.JSONDecodeError:
-        # One retry
+        # retry once if the model didn’t return valid JSON
         raw = call_llm(
             PROMPT_TEMPLATE.format(
                 persona_description=persona_desc,
@@ -142,17 +129,17 @@ def generate_character(world_desc: str, persona_desc: str,
         )
         return json.loads(raw)
 
-# ───────────────────────────── 7.  UI ────────────────────────────────────────
+# ─────────────────────────── 7.  Gradio UI ──────────────────────────────────
 DESCRIPTION = """
-* Generates a character sheet (JSON) from a world + persona.
-* Appearance images via **FLUX-dev**; narrative via **Zephyr-chat** (or Gemma fallback).
-* Personas come from **FinePersonas-Lite**.
+* 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 → Spin the world, then shuffle personas to see very different heroes.
+Tip → Shuffle the world then persona for rapid inspiration.
 """
 
 with gr.Blocks(title="Character Generator", theme="Nymbo/Nymbo_Theme") as demo:
-    gr.Markdown("<h1 style='text-align:center'>🧚‍♀️ Character Generator</h1>")
+    gr.Markdown("<h1 style='text-align:center'>🧝‍♂️ Character Generator</h1>")
     gr.Markdown(DESCRIPTION.strip())
 
     with gr.Row():
@@ -180,4 +167,3 @@ with gr.Blocks(title="Character Generator", theme="Nymbo/Nymbo_Theme") as demo:
     btn_persona.click(random_persona, outputs=[persona_tb])
 
 demo.queue().launch(share=False)
-