Prompt enhancement

app.py

@@ -1,7 +1,11 @@
+from dotenv import load_dotenv
 import gradio as gr
 import numpy as np
 import random
 import json
+import os
+import logging
+import google.generativeai as genai
 
 import spaces
 import torch
@@ -12,6 +16,16 @@ from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMP
 LOADABLE_CLASSES["pikigen"] = LOADABLE_CLASSES["pikigen.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
 ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]
 
+load_dotenv()
+
+# Initialize Gemini API if API key is available
+if os.getenv("GEMINI_API_KEY"):
+    gemini_available = True
+    genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+else:
+    gemini_available = False
+    logging.warning("GEMINI_API_KEY not found in environment variables. Prompt enrichment will not work.")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Freepik/Pikigen-test"
 
@@ -46,17 +60,70 @@ RESOLUTIONS = {
   "square": [
     {"width": 1216, "height": 1216, "label": "1216×1216"},
     {"width": 1024, "height": 1024, "label": "1024×1024"}
-  ],
-  "default": {"width": 1024, "height": 1024, "label": "1024×1024"}
+  ]
 }
 
+# Default resolution
+DEFAULT_RESOLUTION = {"width": 1024, "height": 1024, "label": "1024×1024"}
+
 # Create flattened options for the dropdown
 resolution_options = []
 for category, resolutions in RESOLUTIONS.items():
-    if category != "default":
-        for res in resolutions:
-            resolution_options.append(f"{category.capitalize()} - {res['label']}")
+    resolution_options.append([f"{category.capitalize()}", None])  # Category header
+    for res in resolutions:
+        resolution_options.append([f"  {res['label']}", f"{category}:{res['width']}:{res['height']}"])
 
+def enrich_prompt_with_gemini(prompt, max_tokens=1024):
+    """
+    Enrich a prompt using Google's Gemini API.
+    
+    Args:
+        prompt: The original prompt to enrich
+        max_tokens: Maximum number of tokens for the response
+    
+    Returns:
+        tuple: (enriched_prompt, error_message)
+    """
+    try:
+        if not os.getenv("GEMINI_API_KEY"):
+            return None, "GEMINI_API_KEY not found in environment variables. Please add it to your .env file."
+        
+        model = genai.GenerativeModel('gemini-1.5-flash')
+        
+        enrichment_prompt = f"""
+        You are a prompt enhancer for image generation. 
+        Take the following basic prompt and make it longer, very descriptive, and detailed. 
+        Write the description in a paragraph, avoiding bullet points.
+        
+        Original prompt: {prompt}
+        
+        Enhanced prompt:
+        """
+        
+        response = model.generate_content(enrichment_prompt, generation_config={
+            "max_output_tokens": max_tokens,
+            "temperature": 1,
+        })
+        
+        enriched_prompt = response.text.strip()
+        return enriched_prompt, None
+        
+    except Exception as e:
+        error_message = f"Error enriching prompt: {str(e)}"
+        logging.error(error_message)
+        return None, error_message
+
+# Function to update width and height based on selected resolution
+def update_resolution(resolution_value):
+    """Updates width and height based on selected resolution value"""
+    if not resolution_value:
+        return DEFAULT_RESOLUTION["width"], DEFAULT_RESOLUTION["height"]
+    
+    try:
+        category, width, height = resolution_value.split(":")
+        return int(width), int(height)
+    except:
+        return DEFAULT_RESOLUTION["width"], DEFAULT_RESOLUTION["height"]
 
 @spaces.GPU(duration=120)
 def infer(
@@ -68,15 +135,26 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    use_prompt_enrichment,
     progress=gr.Progress(track_tqdm=True),
 ):
+    enriched_prompt_str = None
+    error_message_str = None
+    generation_prompt = prompt # Default to original prompt
+    
+    if use_prompt_enrichment and gemini_available:
+        enriched_prompt_str, error_message_str = enrich_prompt_with_gemini(prompt)
+        if enriched_prompt_str:
+            generation_prompt = enriched_prompt_str # Use enriched prompt if successful
+        # If enrichment fails, generation_prompt remains the original prompt
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
     image = pipe(
-        prompt=prompt,
+        prompt=generation_prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
@@ -84,33 +162,23 @@ def infer(
         height=height,
         generator=generator,
     ).images[0]
-
-    return image, seed
-
-
-def update_resolution(resolution_option):
-    """Updates width and height based on selected resolution option"""
-    if not resolution_option:
-        # Use default resolution
-        return RESOLUTIONS["default"]["width"], RESOLUTIONS["default"]["height"]
     
-    # Parse the resolution option format: "Category - WidthxHeight"
-    try:
-        category, label = resolution_option.split(" - ")
-        category = category.lower()
-        
-        for res in RESOLUTIONS[category]:
-            if res["label"] == label:
-                return res["width"], res["height"]
-    except:
-        pass
-    
-    # Fallback to default
-    return RESOLUTIONS["default"]["width"], RESOLUTIONS["default"]["height"]
+    # Prepare Gradio updates for the enriched prompt display
+    enriched_prompt_display_update = gr.update(visible=False)
+    enriched_prompt_text_update = gr.update(value="")
+    enrichment_error_update = gr.update(visible=False, value="")
 
+    if enriched_prompt_str:
+        enriched_prompt_display_update = gr.update(visible=True)
+        enriched_prompt_text_update = gr.update(value=enriched_prompt_str)
+    elif error_message_str:
+        enriched_prompt_display_update = gr.update(visible=True)
+        enrichment_error_update = gr.update(visible=True, value=error_message_str)
+
+    return image, seed, enriched_prompt_display_update, enriched_prompt_text_update, enrichment_error_update
 
 examples = [
-    "A photorealistic 3D render of a charming, mischievous young boy, approximately eight years old, possessing the endearingly unusual features of long, floppy donkey ears that droop playfully over his shoulders and a surprisingly small, pink pig nose that twitches slightly.  His eyes, a sparkling, intelligent hazel, are wide with a hint of playful mischief, framed by slightly unruly, sandy-brown hair that falls in tousled waves across his forehead.  He's dressed in a simple, slightly oversized, worn denim shirt and patched-up corduroy trousers, hinting at a life spent playing outdoors. The lighting is soft and natural, casting gentle shadows that highlight the texture of his skin – slightly freckled and sun-kissed, suggesting time spent in the sun.  His expression is one of curious anticipation, his lips slightly parted as if he’s about to speak or perhaps is listening intently. The background is a subtly blurred pastoral scene, perhaps a sun-dappled meadow with wildflowers, enhancing the overall whimsical and slightly surreal nature of the character.  The overall style aims for a blend of realistic rendering with a touch of whimsical cartoonishness, capturing the unique juxtaposition of the boy's human features and his animalistic ears and nose.",
+    "A photorealistic 3D render of a charming, mischievous young boy, approximately eight years old, possessing the endearingly unusual features of long, floppy donkey ears that droop playfully over his shoulders and a surprisingly small, pink pig nose that twitches slightly.  His eyes, a sparkling, intelligent hazel, are wide with a hint of playful mischief, framed by slightly unruly, sandy-brown hair that falls in tousled waves across his forehead.  He's dressed in a simple, slightly oversized, worn denim shirt and patched-up corduroy trousers, hinting at a life spent playing outdoors. The lighting is soft and natural, casting gentle shadows that highlight the texture of his skin – slightly freckled and sun-kissed, suggesting time spent in the sun.  His expression is one of curious anticipation, his lips slightly parted as if he's about to speak or perhaps is listening intently. The background is a subtly blurred pastoral scene, perhaps a sun-dappled meadow with wildflowers, enhancing the overall whimsical and slightly surreal nature of the character.  The overall style aims for a blend of realistic rendering with a touch of whimsical cartoonishness, capturing the unique juxtaposition of the boy's human features and his animalistic ears and nose.",
     "Two white swans with long necks, gracefully swimming in a still body of water. The swans are positioned in a heart shape, with their necks intertwined, creating a romantic and elegant scene. The water is calm and reflective, reflecting the soft, golden light of the setting sun. The background is a blur of soft, golden hues, suggesting a peaceful and serene environment. The image is likely a photograph, captured with a shallow depth of field, which emphasizes the swans and creates a sense of intimacy. The soft lighting and the gentle curves of the swans create a sense of tranquility and beauty. The overall mood of the image is one of love, peace, and serenity.",
     """A watercolor painting of the American flag waving in the wind. The flag is painted in a vibrant red, white, and blue, with the stars in the blue field appearing slightly blurred, creating a sense of motion. The red stripes are painted with a slightly textured brushstroke, giving the flag a realistic and weathered look. The flag is positioned diagonally across the image, with the top left corner extending beyond the frame. The background is a simple white, allowing the flag to be the focal point. Below the flag, in bold red letters, is the text "PRESIDENTS DAY," with "21 FEBRUARY" in blue text above it. Below the text, in black, is "UNITED STATES OF AMERICA." The overall style of the image is patriotic and celebratory, with the watercolor technique adding a touch of artistic flair. The image evokes a sense of pride and national unity, making it a fitting tribute to Presidents Day.""",
     "A captivating photo, shot with a shallow depth of field, of a stunning blonde woman with cascading waves of platinum blonde hair that fall past her shoulders, catching the light. Her eyes, a striking shade of emerald green, are intensely focused on something just off-camera, creating a sense of intrigue.  Sunlight streams softly onto her face, highlighting the delicate curve of her cheekbones and the subtle freckles scattered across her nose. She's wearing a flowing, bohemian-style maxi dress, the fabric a deep sapphire blue that complements her hair and eyes beautifully. The dress is adorned with intricate embroidery along the neckline and sleeves, adding a touch of elegance.  The background is intentionally blurred, suggesting a sun-drenched garden setting with hints of vibrant flowers and lush greenery, drawing the viewer's eye to the woman's captivating features.  The overall mood is serene yet captivating, evoking a feeling of summer warmth and quiet contemplation.  The image should have a natural, slightly ethereal quality, with soft, diffused lighting that enhances her beauty without harsh shadows.",
@@ -119,35 +187,71 @@ examples = [
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 640px;
+    max-width: 1024px;
+}
+.prompt-row > .gr-form {
+    gap: 0.5rem !important; /* Reduce gap between checkbox and button */
+    align-items: center; /* Align items vertically */
 }
 """
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, theme="ParityError/Interstellar") as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # F-lite Text-to-Image Demo")
 
-        with gr.Row():
+        with gr.Row(elem_classes="prompt-row"):
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
                 placeholder="Enter your prompt",
                 container=False,
+                scale=6 # Give prompt more space
+            )
+            
+            use_prompt_enrichment = gr.Checkbox(
+                label="Enrich", 
+                value=True if gemini_available else False,
+                visible=gemini_available, # Hide checkbox if Gemini not available
+                scale=1, # Give checkbox some space
+                min_width=100 # Ensure label isn't cut off
             )
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+            run_button = gr.Button("Run", scale=1, variant="primary", min_width=100)
 
         result = gr.Image(label="Result", show_label=False)
+        
+        # Enriched prompt display (outside Advanced Settings)
+        enriched_prompt_display = gr.Accordion("Enriched Prompt", open=False, visible=False)
+        with enriched_prompt_display:
+            enriched_prompt_text = gr.Textbox(
+                label="Enriched Prompt",
+                interactive=False,
+                lines=8
+            )
+            enrichment_error = gr.Textbox(
+                label="Error",
+                visible=False,
+                interactive=False,
+            )
 
         with gr.Accordion("Advanced Settings", open=False):
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=True,
+            )
+            
+            # Removed checkbox and enriched prompt display from here
+
             with gr.Tabs() as resolution_tabs:
                 with gr.TabItem("Preset Resolutions"):
                     resolution_dropdown = gr.Dropdown(
-                        label="Select Resolution",
-                        choices=[""] + resolution_options,
-                        value="",
-                        allow_custom_value=False,
+                        label="Resolution",
+                        choices=resolution_options,
+                        value="square:1024:1024",
+                        type="value"
                     )
                 
                 with gr.TabItem("Custom Resolution"):
@@ -157,7 +261,7 @@ with gr.Blocks(css=css) as demo:
                             minimum=256,
                             maximum=MAX_IMAGE_SIZE,
                             step=32,
-                            value=RESOLUTIONS["default"]["width"],
+                            value=DEFAULT_RESOLUTION["width"],
                         )
 
                         height = gr.Slider(
@@ -165,25 +269,18 @@ with gr.Blocks(css=css) as demo:
                             minimum=256,
                             maximum=MAX_IMAGE_SIZE,
                             step=32,
-                            value=RESOLUTIONS["default"]["height"],
+                            value=DEFAULT_RESOLUTION["height"],
                         )
 
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=True,
-            )
-
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
                 maximum=MAX_SEED,
                 step=1,
-                value=0,
+                value=42,
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
 
             with gr.Row():
                 guidance_scale = gr.Slider(
@@ -203,15 +300,18 @@ with gr.Blocks(css=css) as demo:
                 )
 
         # Examples should explicitly target only the prompt input
-        gr.Examples(examples=examples, inputs=prompt, example_labels=[ex[:120] + "..." if len(ex) > 120 else ex for ex in examples])
+        max_length = 180
+        gr.Examples(examples=examples, inputs=prompt, example_labels=[ex[:max_length] + "..." if len(ex) > max_length else ex for ex in examples])
     
-    # Update width and height when resolution is selected from dropdown
+    # Update width and height when resolution dropdown changes
     resolution_dropdown.change(
         fn=update_resolution,
-        inputs=[resolution_dropdown],
+        inputs=resolution_dropdown,
         outputs=[width, height]
     )
     
+    # Removed separate function and gr.on for enriched prompt display update
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -224,8 +324,10 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            use_prompt_enrichment,
         ],
-        outputs=[result, seed],
+        # Outputs now include updates for the enriched prompt display components
+        outputs=[result, seed, enriched_prompt_display, enriched_prompt_text, enrichment_error],
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -6,4 +6,6 @@ protobuf
 rich
 sentencepiece
 torch
-transformers
+transformers
+google-generativeai
+dotenv