feat: Integrate Gemini AI for prompt refinement and image generation

- Migrated to google-genai SDK.
- Implemented prompt refinement with gemini-3-pro-preview.
- Integrated image generation using imagen-4.0-generate-001.
- Fixed image display by converting Gemini response to PIL Image.
- Added PNG conversion and friendly filenames (rpg_portrait.png).
- Expanded exotic weapon, armor, and accessory variants in features.yaml.
- Added support for dual accessory selection.
- Implemented "Save Character" and "Load Character" functionality (JSON).
- Added detailed traceback logging for debugging.
- Fixed Gradio theme and CopyButton compatibility issues.

app.py

@@ -2,6 +2,28 @@ import gradio as gr
 import yaml
 import random
 import os
+from dotenv import load_dotenv
+from google import genai
+from google.genai import types
+from PIL import Image
+import io
+import json
+import tempfile
+import traceback
+
+# Load environment variables
+load_dotenv()
+
+# Setup Gemini
+api_key = os.getenv("GEMINI_API_KEY")
+client = None
+gemini_active = False
+if api_key:
+    try:
+        client = genai.Client(api_key=api_key)
+        gemini_active = True
+    except Exception as e:
+        print(f"Error initializing Gemini: {e}")
 
 # Load features from YAML
 def load_features():
@@ -10,59 +32,99 @@ def load_features():
 
 features_data = load_features()
 
-# Define the sequence of features to ensure consistency between logic and UI
-# This MUST match the order components are created in build_ui()
+# Define segments for prompt building
+# (Category, Subcategory, Template Key)
 FEATURE_SEQUENCE = [
-    # Identity
-    ('identity', 'race'), ('identity', 'class'), ('identity', 'gender'), ('identity', 'age'),
-    # Expression & Pose
-    ('expression_pose', 'expression'), ('expression_pose', 'pose'),
-    # Appearance
-    ('appearance', 'hair_color'), ('appearance', 'hair_style'), ('appearance', 'eye_color'), 
-    ('appearance', 'build'), ('appearance', 'skin_tone'), ('appearance', 'distinguishing_feature'),
-    # Equipment
-    ('equipment', 'armor'), ('equipment', 'weapon'), ('equipment', 'accessory'), ('equipment', 'material'),
-    # Environment
-    ('environment', 'background'), ('environment', 'lighting'), ('environment', 'atmosphere'),
-    # Style & Effects
-    ('vfx_style', 'vfx'), ('vfx_style', 'style'), ('vfx_style', 'mood'), ('vfx_style', 'camera'),
-    # Technical
-    ('technical', 'aspect_ratio')
+    ('identity', 'race', 'race'), 
+    ('identity', 'class', 'class'), 
+    ('identity', 'gender', 'gender'), 
+    ('identity', 'age', 'age'),
+    ('expression_pose', 'expression', 'expression'), 
+    ('expression_pose', 'pose', 'pose'),
+    ('appearance', 'hair_color', 'hair_color'), 
+    ('appearance', 'hair_style', 'hair_style'), 
+    ('appearance', 'eye_color', 'eye_color'), 
+    ('appearance', 'build', 'build'), 
+    ('appearance', 'skin_tone', 'skin_tone'), 
+    ('appearance', 'distinguishing_feature', 'distinguishing_feature'),
+    ('equipment', 'armor', 'armor'), 
+    ('equipment', 'weapon', 'weapon'), 
+    ('equipment', 'accessory', 'accessory'), 
+    ('equipment', 'accessory', 'accessory2'), 
+    ('equipment', 'material', 'material'),
+    ('environment', 'background', 'background'), 
+    ('environment', 'lighting', 'lighting'), 
+    ('environment', 'atmosphere', 'atmosphere'),
+    ('vfx_style', 'vfx', 'vfx'), 
+    ('vfx_style', 'style', 'style'), 
+    ('vfx_style', 'mood', 'mood'), 
+    ('vfx_style', 'camera', 'camera'),
+    ('technical', 'aspect_ratio', 'aspect_ratio')
 ]
 
+# Section names for extra info
+SECTIONS = ['Identity', 'Appearance', 'Equipment', 'Environment', 'Style']
+
 def get_detail(category, subcategory, key):
     """Retrieves the detailed description for a given key in a category/subcategory."""
     return features_data.get(category, {}).get(subcategory, {}).get(key, key)
 
 def generate_prompt(*args):
     """
-    Assembles the prompt based on dropdown selections, mapping keys to detailed descriptions.
+    Assembles the prompt based on dropdown selections and extra text info.
     """
     num_features = len(FEATURE_SEQUENCE)
     feature_keys = args[:num_features]
+    extra_infos = args[num_features : num_features + len(SECTIONS)]
 
     template = features_data.get("templates", {}).get("default", "")
     
-    # Build context dictionary dynamically based on FEATURE_SEQUENCE
     context = {}
-    for i, (cat, subcat) in enumerate(FEATURE_SEQUENCE):
+    for i, (cat, subcat, t_key) in enumerate(FEATURE_SEQUENCE):
         key = feature_keys[i]
-        template_var = subcat
-        context[template_var] = get_detail(cat, subcat, key)
+        context[t_key] = get_detail(cat, subcat, key)
+    
+    # Handle multiple accessories dynamically
+    acc_list = []
+    # Identify accessory keys in context
+    for i, (cat, subcat, t_key) in enumerate(FEATURE_SEQUENCE):
+        if subcat == 'accessory' and feature_keys[i] != "None":
+            acc_list.append(context[t_key])
+    
+    if acc_list:
+        context['accessories'] = ", and has " + " as well as ".join(acc_list)
+    else:
+        context['accessories'] = ""
+    
+    # Inject extra info into the context or append to segments
+    # For simplicity, we'll append the extra info to specific segments if present
+    # Mapping section index to template keys
+    # Identity (0-3), Appearance (4-11), Equipment (12-15), Environment (16-18), Style (19-22)
     
-    return template.format(**context)
+    if extra_infos[0]: # Identity
+        context['age'] += f", {extra_infos[0]}"
+    if extra_infos[1]: # Appearance
+        context['distinguishing_feature'] += f", also {extra_infos[1]}"
+    if extra_infos[2]: # Equipment
+        context['accessory'] += f", complemented by {extra_infos[2]}"
+    if extra_infos[3]: # Environment
+        context['atmosphere'] += f", additionally {extra_infos[3]}"
+    if extra_infos[4]: # Style
+        context['camera'] += f", art style notes: {extra_infos[4]}"
+    
+    try:
+        return template.format(**context)
+    except Exception as e:
+        return f"Error building prompt: {e}"
 
 def handle_regeneration(*args):
-    """
-    Checks each randomization checkbox. If checked, picks a random item for the dropdown.
-    Returns the new list of dropdown values.
-    """
+    """Randomizes checkboxes and returns new values for dropdowns."""
     num_features = len(FEATURE_SEQUENCE)
     current_values = list(args[:num_features])
-    checkboxes = args[num_features:]
+    checkboxes = args[num_features : num_features*2]
     
     new_values = []
-    for i, (is_random, (cat, subcat)) in enumerate(zip(checkboxes, FEATURE_SEQUENCE)):
+    for i, (is_random, (cat, subcat, t_key)) in enumerate(zip(checkboxes, FEATURE_SEQUENCE)):
         if is_random:
             choices = list(features_data[cat][subcat].keys())
             new_values.append(random.choice(choices))
@@ -71,13 +133,119 @@ def handle_regeneration(*args):
             
     return new_values
 
+def save_character(*args):
+    """Saves all current UI values to a JSON file."""
+    num_features = len(FEATURE_SEQUENCE)
+    feature_keys = args[:num_features]
+    checkboxes = args[num_features : num_features*2]
+    extra_infos = args[num_features*2 : num_features*2 + len(SECTIONS)]
+    
+    data = {
+        "features": {FEATURE_SEQUENCE[i][2]: feature_keys[i] for i in range(num_features)},
+        "randomization": {FEATURE_SEQUENCE[i][2]: checkboxes[i] for i in range(num_features)},
+        "extra_info": {SECTIONS[i].lower(): extra_infos[i] for i in range(len(SECTIONS))}
+    }
+    
+    # Save to a file with a friendly name in a temp directory
+    temp_dir = tempfile.mkdtemp()
+    path = os.path.join(temp_dir, "rpg_character_config.json")
+    with open(path, 'w') as f:
+        json.dump(data, f, indent=4)
+    return path
+
+def load_character(file):
+    """Loads UI values from a JSON file."""
+    if file is None:
+        return [gr.update()] * (len(FEATURE_SEQUENCE) + len(SECTIONS))
+    
+    try:
+        with open(file.name, 'r') as f:
+            data = json.load(f)
+        
+        updates = []
+        # Update dropdowns
+        for _, _, t_key in FEATURE_SEQUENCE:
+            updates.append(data.get("features", {}).get(t_key, gr.update()))
+        
+        # Note: We don't necessarily update checkboxes unless requested, 
+        # but for a full load, let's just return the values for dropdowns and textboxes first.
+        # Following the pattern of handle_regeneration, we might need to return specific updates.
+        
+        # Update extra info textboxes
+        for section in SECTIONS:
+            updates.append(data.get("extra_info", {}).get(section.lower(), ""))
+            
+        return updates
+    except Exception as e:
+        print(f"Error loading character: {e}")
+        return [gr.update()] * (len(FEATURE_SEQUENCE) + len(SECTIONS))
+
+def refine_with_gemini(prompt):
+    if not gemini_active:
+        return "Gemini API key not found in .env file."
+    
+    system_prompt = (
+        "You are an expert prompt engineer for AI image generators. "
+        "Your task is to take the provided technical prompt and refine it into a more vivid, "
+        "artistic, and detailed description while maintaining all the core features. "
+        "Enhance the vocabulary, lighting, and textures to be professional quality. "
+        "IMPORTANT: RETURN ONLY THE REFINED PROMPT TEXT. DO NOT INCLUDE ANY CONVERSATIONAL FILLER, MARKDOWN CODE BLOCKS, OR EXPLANATIONS."
+    )
+    
+    try:
+        response = client.models.generate_content(
+            model='gemini-3-pro-preview',
+            contents=f"{system_prompt}\n\nOriginal Prompt: {prompt}"
+        )
+        text = response.text.strip()
+        # Remove common markdown code block wrappings if present
+        if text.startswith("```"):
+            lines = text.splitlines()
+            if lines[0].startswith("```"):
+                lines = lines[1:]
+            if lines and lines[-1].startswith("```"):
+                lines = lines[:-1]
+            text = "\n".join(lines).strip()
+        return text
+    except Exception as e:
+        return f"Error refining prompt with Gemini: {e}"
+
+def generate_image_with_gemini(refined_prompt, technical_prompt):
+    if not gemini_active:
+        return None, "Gemini API key not found in .env file."
+    
+    # Priority: Refined Prompt -> Technical Prompt
+    final_prompt = refined_prompt.strip() if refined_prompt and refined_prompt.strip() else technical_prompt.strip()
+    
+    if not final_prompt:
+        return None, "No prompt available for generation."
+    
+    try:
+        # Using the new SDK's generate_images method
+        response = client.models.generate_images(
+            model='imagen-4.0-generate-001',
+            prompt=final_prompt
+        )
+        if response.generated_images:
+            img = Image.open(io.BytesIO(response.generated_images[0].image.image_bytes))
+            # Save as PNG to a temp file for download
+            temp_dir = tempfile.mkdtemp()
+            img_path = os.path.join(temp_dir, "rpg_portrait.png")
+            img.save(img_path, "PNG")
+            return img, img_path, f"Image generated using {'refined' if refined_prompt.strip() else 'technical'} prompt!"
+        return None, None, "Gemini Image generation did not return any images."
+    except Exception as e:
+        traceback.print_exc()
+        return None, f"Image Generation Error: {e}"
+
 def build_ui():
-    with gr.Blocks(title="RPGPortrait Prompt Builder") as demo:
-        gr.Markdown("# 🎨 RPGPortrait Prompt Builder")
-        gr.Markdown("Create highly detailed prompts for advanced AI models like Nano Banana Pro.")
+    with gr.Blocks(title="RPGPortrait Prompt Builder Pro") as demo:
+        gr.Markdown("# 🎨 RPGPortrait Prompt Builder Pro")
+        gr.Markdown("Create professional AI prompts and generate portraits with Gemini integration.")
 
         dropdowns = []
         checkboxes = []
+        extra_texts = []
 
         def create_feature_ui(category, subcategory, label, default_value):
             choices = list(features_data[category][subcategory].keys())
@@ -87,7 +255,6 @@ def build_ui():
             dropdowns.append(dd)
             checkboxes.append(cb)
 
-        # UI Creation order matches FEATURE_SEQUENCE exactly
         with gr.Row():
             with gr.Column():
                 gr.Markdown("### 👤 Identity")
@@ -95,6 +262,8 @@ def build_ui():
                 create_feature_ui('identity', 'class', "Class", "Fighter")
                 create_feature_ui('identity', 'gender', "Gender", "Male")
                 create_feature_ui('identity', 'age', "Age", "Young Adult")
+                extra_id = gr.Textbox(placeholder="Extra Identity details (e.g. lineage, title)", label="Additional Identity Info")
+                extra_texts.append(extra_id)
                 
                 gr.Markdown("### 🎭 Expression & Pose")
                 create_feature_ui('expression_pose', 'expression', "Expression", "Determined")
@@ -108,20 +277,27 @@ def build_ui():
                 create_feature_ui('appearance', 'build', "Build", "Average")
                 create_feature_ui('appearance', 'skin_tone', "Skin Tone", "Fair")
                 create_feature_ui('appearance', 'distinguishing_feature', "Distinguishing Feature", "None")
+                extra_app = gr.Textbox(placeholder="Extra Appearance details (e.g. tattoos, scars)", label="Additional Appearance Info")
+                extra_texts.append(extra_app)
 
         with gr.Row():
             with gr.Column():
                 gr.Markdown("### ⚔️ Equipment")
                 create_feature_ui('equipment', 'armor', "Armor/Clothing", "Travelers Clothes")
                 create_feature_ui('equipment', 'weapon', "Primary Weapon", "Longsword")
-                create_feature_ui('equipment', 'accessory', "Accessory", "None")
+                create_feature_ui('equipment', 'accessory', "Accessory 1", "None")
+                create_feature_ui('equipment', 'accessory', "Accessory 2", "None")
                 create_feature_ui('equipment', 'material', "Material Detail", "Weathered")
+                extra_eq = gr.Textbox(placeholder="Extra Equipment details (e.g. weapon enchantments)", label="Additional Equipment Info")
+                extra_texts.append(extra_eq)
 
             with gr.Column():
                 gr.Markdown("### 🌍 Environment")
                 create_feature_ui('environment', 'background', "Background", "Forest")
                 create_feature_ui('environment', 'lighting', "Lighting", "Natural Sunlight")
                 create_feature_ui('environment', 'atmosphere', "Atmosphere", "Clear")
+                extra_env = gr.Textbox(placeholder="Extra Environment details (e.g. time of day, weather)", label="Additional Environment Info")
+                extra_texts.append(extra_env)
 
         with gr.Row():
             with gr.Column():
@@ -130,6 +306,8 @@ def build_ui():
                 create_feature_ui('vfx_style', 'style', "Art Style", "Digital Illustration")
                 create_feature_ui('vfx_style', 'mood', "Mood", "Heroic")
                 create_feature_ui('vfx_style', 'camera', "Camera Angle", "Bust")
+                extra_sty = gr.Textbox(placeholder="Extra Style details (e.g. specific artists, colors)", label="Additional Style Info")
+                extra_texts.append(extra_sty)
 
             with gr.Column():
                 gr.Markdown("### ⚙️ Technical")
@@ -138,37 +316,79 @@ def build_ui():
         gr.Markdown("---")
         
         with gr.Row():
-            regenerate_btn = gr.Button("✨ Regenerate Random Features", variant="primary")
+            regenerate_btn = gr.Button("✨ Randomize Features", variant="secondary")
+            save_btn = gr.Button("💾 Save Character", variant="secondary")
+            load_btn = gr.UploadButton("📂 Load Character", file_types=[".json"], variant="secondary")
+            refine_btn = gr.Button("🧠 Refine with Gemini", variant="primary")
+            gen_img_btn = gr.Button("🖼️ Generate Image (Gemini)", variant="primary")
         
         with gr.Row():
-            prompt_output = gr.Textbox(
-                label="Generated Prompt for Nano Banana Pro",
-                lines=6,
-                interactive=False
-            )
+            with gr.Column(scale=2):
+                prompt_output = gr.Textbox(label="Generated Technical Prompt", lines=5, interactive=False)
+                refined_output = gr.Textbox(label="Gemini Refined Prompt", lines=5, interactive=True)
+            with gr.Column(scale=1):
+                image_output = gr.Image(label="Gemini Generated Portrait")
+                download_img_btn = gr.DownloadButton("📥 Download Portrait (PNG)", variant="secondary", visible=False)
+                download_file = gr.File(label="Saved Character JSON", visible=False)
+                status_msg = gr.Markdown("")
 
-        all_inputs = dropdowns + checkboxes
-
-        # Set up event listeners for automatic updates on dropdown change
-        for dd in dropdowns:
-            dd.change(fn=generate_prompt, inputs=dropdowns, outputs=prompt_output)
+        all_input_components = dropdowns + extra_texts
+        
+        # Automatic prompt update on any input change
+        for comp in all_input_components:
+            comp.change(fn=generate_prompt, inputs=all_input_components, outputs=prompt_output)
         
         # Regenerate button logic
         regenerate_btn.click(
             fn=handle_regeneration,
-            inputs=all_inputs,
+            inputs=dropdowns + checkboxes,
             outputs=dropdowns
         ).then(
+            # Trigger prompt update after dropdowns change
             fn=generate_prompt,
-            inputs=dropdowns,
+            inputs=all_input_components,
             outputs=prompt_output
         )
+
+        # Gemini Refinement
+        refine_btn.click(
+            fn=refine_with_gemini,
+            inputs=prompt_output,
+            outputs=refined_output
+        )
+
+        # Gemini Image Generation
+        gen_img_btn.click(
+            fn=generate_image_with_gemini,
+            inputs=[refined_output, prompt_output], 
+            outputs=[image_output, download_img_btn, status_msg]
+        ).then(
+            fn=lambda x: gr.update(value=x, visible=True) if x else gr.update(visible=False),
+            inputs=download_img_btn,
+            outputs=download_img_btn
+        )
+
+        # Save/Load Logic
+        save_btn.click(
+            fn=save_character,
+            inputs=dropdowns + checkboxes + extra_texts,
+            outputs=download_file
+        ).then(
+            fn=lambda: gr.update(visible=True),
+            outputs=download_file
+        )
+
+        load_btn.upload(
+            fn=load_character,
+            inputs=load_btn,
+            outputs=dropdowns + extra_texts
+        )
         
-        # Initialize the prompt on load
-        demo.load(fn=generate_prompt, inputs=dropdowns, outputs=prompt_output)
+        # Initialize
+        demo.load(fn=generate_prompt, inputs=all_input_components, outputs=prompt_output)
 
     return demo
 
 if __name__ == "__main__":
     demo = build_ui()
-    demo.launch()
+    demo.launch(theme=gr.themes.Soft())

features.yaml

@@ -105,6 +105,9 @@ equipment:
     Chainmail: "clad in a suit of shimmering, interlocking chainmail"
     Travelers Clothes: "dressed in practical, durable travelers' garments"
     Elegant Robes: "wearing flowing, silk-threaded elegant robes with intricate embroidery"
+    Bone Armor: "wearing intimidating, primal armor crafted from the bones of prehistoric beasts"
+    Arcane Plate: "clad in futuristic, bioluminescent plate armor etched with glowing circuitry"
+    Leaf-weave Armor: "attired in flexible armor made of magically hardened, vibrant-green leaves"
     Ragged Tunic: "attired in a simple, worn, and slightly ragged tunic"
   weapon:
     Longsword: "confidently wielding a gleaming steel longsword in a two-handed grip"
@@ -112,6 +115,14 @@ equipment:
     Daggers: "gripping a pair of obsidian daggers, ready for a swift strike"
     Bow and Quiver: "carrying a finely crafted longbow with a full quiver of arrows slung over the shoulder"
     Greataxe: "resting a massive, notched greataxe upon one shoulder"
+    Katana: "holding a slender, razor-sharp katana with a traditional silken hilt-wrap"
+    Whip-blade: "wielding a segmented, metallic whip-blade that can transition between sword and whip"
+    Twin Scimitars: "deftly handling a pair of curved, engraved scimitars"
+    Meteor Hammer: "swinging a heavy, spiked metal sphere attached to a long, sturdy chain"
+    Primitive Club: "gripping a heavy, notched primitive club made of dark, weathered wood"
+    Rapier: "wielding a slender, elegant rapier with an intricate basket hilt"
+    Morning Star: "holding a lethal, spiked morning star, its heavy head gleaming with menace"
+    Sword and Shield: "wielding a balanced shortsword in one hand and a sturdy, emblazoned heater shield in the other"
     Shield: "defensively holding a heavy heater shield emblazoned with a crest"
     None: "with hands held open and empty"
   accessory:
@@ -119,6 +130,12 @@ equipment:
     Hood: "a deep hood casting shadows over the face"
     Amulet: "a glowing amulet hanging from a silver chain"
     Belt with Pouches: "a sturdy leather belt laden with various pouches and tools"
+    Feathered Hat: "wearing a wide-brimmed hat adorned with a long, colorful feather"
+    Silk Scarf: "wrapped in a long, vibrant silk scarf that catches the breeze"
+    Adventurer's Backpack: "shouldering a heavy leather backpack laden with supplies and rope"
+    Oil Lantern: "carrying a flickering oil lantern that casts a warm, localized glow"
+    Ornate Mask: "wearing a decorative gold and velvet masquerade mask"
+    Jeweled Circlet: "adorned with a delicate jeweled circlet that rests upon the brow"
     Eyepatch: "a leather eyepatch covering one eye"
     None: "no visible accessories"
   material:
@@ -193,4 +210,4 @@ technical:
     "2:3": "--ar 2:3"
 
 templates:
-  default: "A {mood} {style} portrait of {race} {class}, {gender}, {age}. {build} build, {hair_style} {hair_color} hair, {eye_color} eyes, with {distinguishing_feature}. The character is {expression} while {pose}, {armor}, {weapon}, and has {accessory}. All equipment is {material}. The setting is a {background}, {lighting}, {atmosphere}. {vfx}. {camera}. Highly detailed, masterpiece, 8k resolution. {aspect_ratio}"
+  default: "A {mood} {style} portrait of {race} {class}, {gender}, {age}. {build} build, {hair_style} {hair_color} hair, {eye_color} eyes, with {distinguishing_feature}. The character is {expression} while {pose}, {armor}, {weapon}{accessories}. All equipment is {material}. The setting is a {background}, {lighting}, {atmosphere}. {vfx}. {camera}. Highly detailed, 8k resolution. {aspect_ratio}"

requirements.txt

@@ -1,2 +1,4 @@
 gradio
 PyYAML
+python-dotenv
+google-genai