Enhance character generation by expanding CharacterDescription model and updating prompt template for JSON output

app.py

@@ -6,11 +6,23 @@ from diffusers import FluxPipeline, AutoencoderKL
 from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
 import spaces
 from google import genai
+from pydantic import BaseModel
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Configure Gemini API
+# Initialize Google Gemini API client
 llm_client = genai.Client()
+class CharacterDescription(BaseModel):
+    name: str
+    background: str
+    appearance: str
+    personality: str
+    skills_and_abilities: str
+    goals: str
+    conflicts: str
+    backstory: str
+    current_situation: str
+    spoken_lines: list[str]
 
 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)
@@ -22,22 +34,11 @@ pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_ret
 
 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 '```'.
+prompt_template = """Generate a character with this persona description:
+{persona_description}
+---
+In a world with this description:
+{world_description}
 """
 
 world_description_prompt = "Generate a unique and random world description (Don't Write anything else except the world description)."
@@ -74,13 +75,9 @@ def generate_character(world_description, persona_description, progress=gr.Progr
         world_description=world_description
     )
     response = llm_client.models.generate_content(
-        model="gemini-2.5-flash", contents=prompt_content
+        model="gemini-2.5-flash", contents=prompt_content, config={"response_mime_type": "application/json", "response_schema": CharacterDescription}
     )
-    result = response.text
-    if "</think>" in result:
-        result = result[result.index("</think>")+len("</think>"):].strip()
-    output = json.loads(result)
-    return output
+    return response.parsed
 
 app_description = """
 - This app generates a character in JSON format based on a persona description and a world description.