Update app.py

app.py

@@ -1,48 +1,63 @@
-# app.py
 from datasets import load_dataset
 import gradio as gr
 from gradio_client import Client
-import json, os, random, torch
+import json, os, random, torch, spaces
 from diffusers import FluxPipeline, AutoencoderKL
 from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
-import spaces
 
-# ─────────────────────────────── 1. Device ────────────────────────────────
+# ───────────────────────────── 1.  Device ────────────────────────────────────
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ─────────────────────── 2. Image / FLUX pipeline ─────────────────────────
+# ─────────────────────── 2.  Image / FLUX pipeline ───────────────────────────
 pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.bfloat16
+    "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
+    "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)
 )
 
-# ───────────────────────── 3. LLM (Zephyr-chat) ───────────────────────────
-llm_client = Client("HuggingFaceH4/zephyr-chat")     # public Space
-CHAT_API   = llm_client.view_api()[0]["api_name"]    # e.g. "/chat"
+# ───────────────────────── 3.  LLM client (robust) ───────────────────────────
+def _first_working_client(candidates: list[str]) -> Client:
+    """
+    Try a list of Space URLs / repo-ids, return the first that gives a JSON config.
+    """
+    for src in candidates:
+        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}")
+            return c
+        except Exception as e:
+            print(f"[warn] {src} not usable → {e}")
+    raise RuntimeError("No usable LLM 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_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 Zephyr, a helpful and creative assistant.",
+    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:
     """
-    Robust wrapper around the Zephyr chat Space.
-    Falls back to '...' on any error so the Gradio UI never crashes.
+    Unified chat wrapper – works for both Zephyr and Gemma Spaces.
     """
     history = history or []
     try:
-        # Zephyr-chat expects: prompt, system_prompt, history, temperature, top_p, max_new_tokens
         result = llm_client.predict(
             user_prompt,
             system_prompt,
@@ -52,21 +67,22 @@ def call_llm(
             max_tokens,
             api_name=CHAT_API,
         )
-        # Some Spaces return a plain string, others return the old tuple format.
-        return result.strip() if isinstance(result, str) else result[1][0][-1].strip()
+        # 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"[LLM error] {e}")
-        return "..."
+        print(f"[error] LLM call failed → {e}")
+        return "…"
 
-# ───────────────────────── 4. Persona dataset ────────────────────────────
+# ───────────────────────── 4.  Persona dataset ───────────────────────────────
 ds = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")
 
-def get_random_persona_description() -> str:
-    idx = random.randint(0, len(ds) - 1)
-    return ds[idx]["persona"]
+def random_persona() -> str:
+    return ds[random.randint(0, len(ds) - 1)]["persona"]
 
-# ─────────────────────────── 5. Prompts  ─────────────────────────────────
-prompt_template = """Generate a character with this persona description:
+# ─────────────────────────── 5.  Prompt templates ───────────────────────────
+PROMPT_TEMPLATE = """Generate a character with this persona description:
 
 {persona_description}
 
@@ -74,25 +90,25 @@ In a world with this description:
 
 {world_description}
 
-Write the character in JSON format with these keys:
-name, background, appearance, personality, skills_and_abilities, goals,
-conflicts, backstory, current_situation, spoken_lines (list of strings).
+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 **only** the JSON (no markdown, no fencing)."""
+Respond with JSON only (no markdown)."""
 
-world_description_prompt = (
+WORLD_PROMPT = (
     "Invent a short, unique and vivid world description. "
     "Respond with the description only."
 )
 
-# ─────────────────────── 6. Gradio helper funcs ─────────────────────────
-def get_random_world_description() -> str:
-    return call_llm(world_description_prompt)
+# ─────────────────────── 6.  Helper functions ───────────────────────────────
+def random_world() -> str:
+    return call_llm(WORLD_PROMPT)
 
 @spaces.GPU(duration=75)
 def infer_flux(character_json):
-    """Stream intermediate images while FLUX denoises."""
-    for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=character_json["appearance"],
         guidance_scale=3.5,
         num_inference_steps=28,
@@ -102,82 +118,66 @@ def infer_flux(character_json):
         output_type="pil",
         good_vae=good_vae,
     ):
-        yield image
+        yield img
 
-def generate_character(world_description: str,
-                       persona_description: str,
+def generate_character(world_desc: str, persona_desc: str,
                        progress=gr.Progress(track_tqdm=True)):
     raw = call_llm(
-        prompt_template.format(
-            persona_description=persona_description,
-            world_description=world_description,
+        PROMPT_TEMPLATE.format(
+            persona_description=persona_desc,
+            world_description=world_desc,
         ),
         max_tokens=1024,
     )
     try:
         return json.loads(raw)
     except json.JSONDecodeError:
-        # One retry if the LLM hallucinated
+        # One retry
         raw = call_llm(
-            prompt_template.format(
-                persona_description=persona_description,
-                world_description=world_description,
+            PROMPT_TEMPLATE.format(
+                persona_description=persona_desc,
+                world_description=world_desc,
             ),
             max_tokens=1024,
         )
         return json.loads(raw)
 
-# ───────────────────────────── 7. UI ─────────────────────────────────────
-app_description = """
-- Generates a character profile (JSON) from a world + persona description.
-- **Appearance** images come from [FLUX-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev).
-- **Back-stories** come from [Zephyr-7B-β](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta).
-- Personas are sampled from [FinePersonas-Lite](https://huggingface.co/datasets/MohamedRashad/FinePersonas-Lite).
+# ───────────────────────────── 7.  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**.
 
-Tip → Write or randomise a world, then spin the persona box to see how the same
-world shapes different heroes.
+Tip → Spin the world, then shuffle personas to see very different heroes.
 """
 
 with gr.Blocks(title="Character Generator", theme="Nymbo/Nymbo_Theme") as demo:
-    gr.Markdown("<h1 style='text-align:center'>🧙‍♀️ Character Generator</h1>")
-    gr.Markdown(app_description.strip())
+    gr.Markdown("<h1 style='text-align:center'>🧚‍♀️ Character Generator</h1>")
+    gr.Markdown(DESCRIPTION.strip())
 
     with gr.Row():
-        world_description = gr.Textbox(label="World Description", lines=10, scale=4)
-        persona_description = gr.Textbox(
-            label="Persona Description",
-            value=get_random_persona_description(),
-            lines=10,
-            scale=1,
+        world_tb = gr.Textbox(label="World Description", lines=10, scale=4)
+        persona_tb = gr.Textbox(
+            label="Persona Description", value=random_persona(), lines=10, scale=1
         )
 
     with gr.Row():
-        random_world_btn   = gr.Button("🔄 Random World",   variant="secondary")
-        submit_btn         = gr.Button("✨ Generate Character", variant="primary", scale=5)
-        random_persona_btn = gr.Button("🔄 Random Persona", variant="secondary")
+        btn_world   = gr.Button("🔄 Random World",   variant="secondary")
+        btn_generate = gr.Button("✨ Generate Character", variant="primary", scale=5)
+        btn_persona = gr.Button("🔄 Random Persona", variant="secondary")
 
     with gr.Row():
-        character_image = gr.Image(label="Character Image")
-        character_json  = gr.JSON(label="Character Description")
-
-    # Hooks
-    submit_btn.click(
-        generate_character,
-        inputs=[world_description, persona_description],
-        outputs=[character_json],
+        img_out  = gr.Image(label="Character Image")
+        json_out = gr.JSON(label="Character Description")
+
+    btn_generate.click(
+        generate_character, [world_tb, persona_tb], [json_out]
     ).then(
-        infer_flux,
-        inputs=[character_json],
-        outputs=[character_image],
+        infer_flux, [json_out], [img_out]
     )
 
-    random_world_btn.click(
-        get_random_world_description,
-        outputs=[world_description],
-    )
-    random_persona_btn.click(
-        get_random_persona_description,
-        outputs=[persona_description],
-    )
+    btn_world.click(random_world, outputs=[world_tb])
+    btn_persona.click(random_persona, outputs=[persona_tb])
 
 demo.queue().launch(share=False)
+