Update app.py

app.py

@@ -1,62 +1,97 @@
+# app.py
 from datasets import load_dataset
 import gradio as gr
 from gradio_client import Client
-import json
-import torch
+import json, os, random, torch
 from diffusers import FluxPipeline, AutoencoderKL
 from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
 import spaces
 
+# ─────────────────────────────── 1. Device ────────────────────────────────
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-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.enable_sequential_cpu_offload()
-# pipe.vae.enable_slicing()
-# pipe.vae.enable_tiling()
-# pipe.to(torch.float16)
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
-
-llm_client = Client("HuggingFaceH4/zephyr-7b-beta") # or any other public gradio space/model
-# t2i_client = Client("black-forest-labs/FLUX.1-dev")
-# t2i_client = Client("black-forest-labs/FLUX.1-schnell")
+# ─────────────────────── 2. Image / FLUX 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)
+)
+
+# ───────────────────────── 3. LLM (Zephyr-chat) ───────────────────────────
+llm_client = Client("HuggingFaceH4/zephyr-chat")     # public Space
+CHAT_API   = llm_client.view_api()[0]["api_name"]    # e.g. "/chat"
+
+def call_llm(
+    user_prompt: str,
+    system_prompt: str = "You are Zephyr, a helpful and 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.
+    """
+    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,
+            history,
+            temperature,
+            top_p,
+            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()
+    except Exception as e:
+        print(f"[LLM error] {e}")
+        return "..."
 
+# ───────────────────────── 4. Persona dataset ────────────────────────────
 ds = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")
 
-prompt_template = """Generate a character with this persona description: {persona_description}
-In a world with this description: {world_description}
-
-Write the character in json format with the following fields:
-- name: The name of the character
-- background: The background of the character
-- appearance: The appearance of the character
-- personality: The personality of the character
-- skills_and_abilities: The skills and abilities of the character
-- goals: The goals of the character
-- conflicts: The conflicts of the character
-- backstory: The backstory of the character
-- current_situation: The current situation of the character
-- spoken_lines: The spoken lines of the character (list of strings)
-
-Don't write anything else except the character description in json format and don't include '```'.
-"""
+def get_random_persona_description() -> str:
+    idx = random.randint(0, len(ds) - 1)
+    return ds[idx]["persona"]
 
-world_description_prompt = "Generate a unique and random world description (Don't Write anything else except the world description)."
+# ─────────────────────────── 5. Prompts  ─────────────────────────────────
+prompt_template = """Generate a character with this persona description:
 
-def get_random_world_description():
-    result = llm_client.predict(
-        query=world_description_prompt,
-        history=[],
-        system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-        api_name="/model_chat",
-    )
-    return result[1][0][-1]
+{persona_description}
+
+In a world with this description:
+
+{world_description}
 
-def get_random_persona_description():
-    return ds.shuffle().select([100])[0]["persona"]
+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).
+
+Respond with **only** the JSON (no markdown, no fencing)."""
+
+world_description_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)
 
 @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(
         prompt=character_json["appearance"],
         guidance_scale=3.5,
@@ -69,59 +104,80 @@ def infer_flux(character_json):
     ):
         yield image
 
-def generate_character(world_description, persona_description, progress=gr.Progress(track_tqdm=True)):
-    result = llm_client.predict(
-        query=prompt_template.format(
-            persona_description=persona_description, world_description=world_description
+def generate_character(world_description: str,
+                       persona_description: str,
+                       progress=gr.Progress(track_tqdm=True)):
+    raw = call_llm(
+        prompt_template.format(
+            persona_description=persona_description,
+            world_description=world_description,
         ),
-        history=[],
-        system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-        api_name="/model_chat",
+        max_tokens=1024,
     )
-    output = json.loads(result[1][0][-1])
-    return output
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        # One retry if the LLM hallucinated
+        raw = call_llm(
+            prompt_template.format(
+                persona_description=persona_description,
+                world_description=world_description,
+            ),
+            max_tokens=1024,
+        )
+        return json.loads(raw)
 
+# ───────────────────────────── 7. UI ─────────────────────────────────────
 app_description = """
-- This app generates a character in JSON format based on a persona description and a world description.
-- The character's appearance is generated using [FLUX-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) and the character description is generated using [Qwen2.5-72B-Instruct](https://huggingface.co/Qwen/Qwen2.5-72B-Instruct).
-- The persona description is randomly selected from the [FinePersonas-Lite](https://huggingface.co/datasets/MohamedRashad/FinePersonas-Lite) dataset.
+- 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).
 
-**Note:** I recommend starting with the world description (you can write one or loop over randomly generated ones) and then try different persona descriptions to generate interesting characters for the world you created.
+Tip → Write or randomise a world, then spin the persona box to see how the same
+world shapes different heroes.
 """
 
-with gr.Blocks(title="Character Generator", theme="Nymbo/Nymbo_Theme") as app:
-    with gr.Column():
-        gr.HTML("<center><h1>Character Generator</h1></center>")
-        gr.Markdown(app_description.strip())
-        with gr.Column():
-            with gr.Row():
-                world_description = gr.Textbox(lines=10, label="World Description", scale=4)
-                persona_description = gr.Textbox(lines=10, label="Persona Description", value=get_random_persona_description(), scale=1)
-            with gr.Row():
-                random_world_button = gr.Button(value="Get Random World Description", variant="secondary", scale=1)
-                submit_button = gr.Button(value="Generate Interesting Character!", variant="primary", scale=5)
-                random_persona_button = gr.Button(value="Get Random Persona Description", variant="secondary", scale=1)
-        with gr.Row():
-            character_image = gr.Image(label="Character Image")
-            character_json = gr.JSON(label="Character Description")
-
-        examples = gr.Examples(
-            [
-                "In a world where magic is real and dragons roam the skies, a group of adventurers set out to find the legendary sword of the dragon king.",
-                "Welcome to Aethoria, a vast and mysterious realm where the laws of physics bend to the will of ancient magic. This world is comprised of countless floating islands suspended in an endless sky, each one a unique ecosystem teeming with life and secrets. The islands of Aethoria range from lush, verdant jungles to barren, crystalline deserts. Some are no larger than a city block, while others span hundreds of miles. Connecting these disparate landmasses are shimmering bridges of pure energy, and those brave enough to venture off the beaten path can find hidden portals that instantly transport them across great distances. Aethoria's inhabitants are as diverse as its landscapes. Humans coexist with ethereal beings of light, rock-skinned giants, and shapeshifting creatures that defy classification. Ancient ruins dot the islands, hinting at long-lost civilizations and forgotten technologies that blur the line between science and sorcery. The world is powered by Aether, a mystical substance that flows through everything. Those who can harness its power become formidable mages, capable of manipulating reality itself. However, Aether is a finite resource, and its scarcity has led to conflicts between the various factions vying for control. In the skies between the islands, magnificent airships sail on currents of magic, facilitating trade and exploration. Pirates and sky raiders lurk in the cloudy depths, always on the lookout for unsuspecting prey. Deep beneath the floating lands lies the Undervoid, a dark and treacherous realm filled with nightmarish creatures and untold riches. Only the bravest adventurers dare to plumb its depths, and fewer still return to tell the tale. As an ever-present threat, the Chaos Storms rage at the edges of the known world, threatening to consume everything in their path. It falls to the heroes of Aethoria to uncover the secrets of their world and find a way to push back the encroaching darkness before it's too late. In Aethoria, every island holds a story, every creature has a secret, and every adventure could change the fate of this wondrous, imperiled world.",
-                "In a world from my imagination, there is a city called 'Orakis'. floating in the sky on pillars of pure light. The walls of the city are made of crystal glass, constantly reflecting the colors of dawn and dusk, giving it an eternal celestial glow. The buildings breathe and change their shapes according to the seasons—they grow in spring, strengthen in summer, and begin to fade in autumn until they become mist in winter.",
-            ],
-            world_description,
+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())
+
+    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,
         )
 
-    submit_button.click(
-        generate_character, [world_description, persona_description], outputs=[character_json]
-    ).then(fn=infer_flux, inputs=[character_json], outputs=[character_image])
-    random_world_button.click(
-        get_random_world_description, outputs=[world_description]
+    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")
+
+    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],
+    ).then(
+        infer_flux,
+        inputs=[character_json],
+        outputs=[character_image],
+    )
+
+    random_world_btn.click(
+        get_random_world_description,
+        outputs=[world_description],
     )
-    random_persona_button.click(
-        get_random_persona_description, outputs=[persona_description]
+    random_persona_btn.click(
+        get_random_persona_description,
+        outputs=[persona_description],
     )
 
-app.queue().launch(share=False)
+demo.queue().launch(share=False)