Update app.py

app.py

@@ -37,7 +37,6 @@ FLUX_SPECS = {
 
 # --- Image Processing Pipeline ---
 def preprocess_image(img):
-    """Convert and enhance uploaded images"""
     try:
         if isinstance(img, str):  # Handle file paths
             img = Image.open(img)
@@ -46,44 +45,49 @@ def preprocess_image(img):
         img = img.filter(ImageFilter.SHARPEN)
         return img
     except Exception as e:
-        raise ValueError(f"Image processing error: {str(e)}")
+        raise ValueError(f"🔴 Image processing failed: {str(e)}")
 
 # --- Core Generation Engine ---
 def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi):
     try:
-        # Validate inputs
+        # Step 1: Input Validation
         if not image:
-            return {"error": "Please upload an image"}
-            
+            return {"error": "⚠️ Please upload an image."}, None, None
+
         api_key = api_key or GEMINI_KEY
         if not api_key:
-            return {"error": "API key required - set in env (GEMINI_KEY) or input field"}
+            return {"error": "🔑 API key required - set in env (GEMINI_KEY) or input field."}, None, None
 
-        # Import and configure Gemini only when needed
+        # Step 2: Gemini Setup
         try:
             import google.generativeai as genai
             genai.configure(api_key=api_key)
             model = genai.GenerativeModel("gemini-1.5-pro")
         except ImportError:
-            return {"error": "Failed to import google.generativeai. Install with: pip install google-generativeai"}
+            return {"error": "🚫 Failed to import google.generativeai. Install with: pip install google-generativeai"}, None, None
         except Exception as e:
             if "authentication" in str(e).lower():
-                return {"error": "Invalid API key or authentication error"}
+                return {"error": "🔐 Invalid API key or authentication error"}, None, None
             else:
-                return {"error": f"API initialization error: {str(e)}"}
+                return {"error": f"⚙️ API initialization error: {str(e)}"}, None, None
 
-        # Process image with timeout protection
-        start_time = time.time()
-        img = preprocess_image(image)
-        img_bytes = io.BytesIO()
-        img.save(img_bytes, format="PNG")
-        img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
+        # Step 3: Preprocess Image
+        try:
+            img = preprocess_image(image)
+            img_bytes = io.BytesIO()
+            img.save(img_bytes, format="PNG")
+            img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
+        except Exception as e:
+            return {"error": f"🖼️ Image preparation failed: {str(e)}"}, None, None
 
-        # Build instruction
-        instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
-        instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
+        # Step 4: Build Instruction Prompt
+        try:
+            instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
+            instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
+        except KeyError:
+            return {"error": "🛠️ Invalid style selected. Please choose from available options."}, None, None
 
-        # Generate prompt with timeout protection
+        # Step 5: Call Gemini API
         try:
             response = model.generate_content(
                 contents=[instruction, {"mime_type": "image/png", "data": img_b64}],
@@ -91,121 +95,113 @@ def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color
             )
             raw_prompt = response.text
         except Exception as e:
-            return {"error": f"Generation failed: {str(e)}"}
+            return {"error": f"🤖 Prompt generation failed: {str(e)}"}, None, None
 
-        # Simple quality validation
+        # Step 6: Quality & Token Stats
         validation = {"score": 8, "issues": [], "suggestions": []}
-        
-        # Token tracking
-        input_tokens = len(img_b64) // 4  # Approximate base64 token count
+        input_tokens = len(img_b64) // 4
         output_tokens = len(raw_prompt.split())
-        
+
         return {
             "prompt": raw_prompt,
             "validation": validation,
             "stats": {"input": input_tokens, "output": output_tokens}
-        }
+        }, validation, {"tokens": f"Input: {input_tokens}, Output: {output_tokens}"}
 
     except Exception as e:
         traceback.print_exc()
-        return {"error": str(e)}
+        return {"error": f"💥 Unexpected error: {str(e)}"}, None, None
 
-# --- UI Response Formatting ---
-def format_generation_response(result):
-    """Format the response from generate_prompt for the UI"""
-    if "error" in result:
-        return result["error"], None, None
-    else:
-        return result.get("prompt", ""), result.get("validation", {}), result.get("stats", {})
 
-# Modern copy function using Gradio's JavaScript API
+# --- Modern Copy Function ---
 def copy_text(text):
-    return gr.update(value=text), f"✓ Copied: '{text[:20]}...'", gr.Button.update(variant="secondary")
+    return text, "✓ Copied to clipboard!", gr.Button.update(variant="secondary")
+
 
 # --- Main Interface ---
 def build_interface():
+    global STYLE_INSTRUCTIONS  # Ensure access to global dict
+
     with gr.Blocks(title="Flux Pro Generator", theme="soft") as app:
-        # Header
         gr.Markdown("# 🎨 Flux Pro Prompt Generator")
-        gr.Markdown("Generate optimized design prompts from images using Google's Gemini")
-        
-        # Security Section
-        api_key = gr.Textbox(
-            label="🔑 Gemini API Key",
-            value=GEMINI_KEY,
-            type="password",
-            info="Set GEMINI_KEY environment variable for production"
-        )
+        gr.Markdown("Generate optimized design prompts from images using Google's Gemini.")
 
-        # Main Workflow
         with gr.Row():
             with gr.Column(scale=1):
+                api_key = gr.Textbox(
+                    label="🔑 Gemini API Key",
+                    value=GEMINI_KEY,
+                    type="password",
+                    info="Set GEMINI_KEY environment variable for production."
+                )
+
                 img_input = gr.Image(
                     label="🖼️ Upload Design",
                     type="pil",
                     sources=["upload"],
                     interactive=True
                 )
+
                 style = gr.Dropdown(
                     list(STYLE_INSTRUCTIONS.keys()),
                     value="General",
                     label="🎨 Target Style"
                 )
-                
-                # Advanced Settings
+
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
-                    creativity = gr.Slider(0.0, 1.0, 0.7, label="Creativity Level")
+                    creativity = gr.Slider(0.0, 1.0, 0.7, label="🧠 Creativity Level")
                     neg_prompt = gr.Textbox(label="🚫 Negative Prompts", placeholder="What to avoid")
                     aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
                     color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], value="RGB", label="Color Mode")
                     dpi = gr.Dropdown([str(d) for d in FLUX_SPECS["dpi_options"]], value="300", label="Output DPI")
-                
+
                 gen_btn = gr.Button("✨ Generate Prompt", variant="primary")
 
             with gr.Column(scale=2):
+                status_msg = gr.Textbox(label="📢 Status", interactive=False)
+
                 prompt_output = gr.Textbox(
                     label="📝 Optimized Prompt",
                     lines=8,
                     interactive=True,
-                    show_copy_button=True  # Modern Gradio has built-in copy button
+                    show_copy_button=True
                 )
-                status_msg = gr.Textbox(label="Status", visible=True)
+
                 with gr.Row():
                     copy_btn = gr.Button("📋 Copy to Clipboard", variant="secondary")
-                
-                quality_report = gr.JSON(
-                    label="🔍 Quality Report",
-                    visible=True
-                )
-                token_stats = gr.JSON(
-                    label="🧮 Token Usage",
-                    visible=True
-                )
-        
+
+                quality_report = gr.JSON(label="🔍 Quality Report")
+                token_stats = gr.JSON(label="🧮 Token Usage")
+
         # Event Handling
         gen_btn.click(
             fn=generate_prompt,
             inputs=[
-                img_input, api_key, style, creativity, 
+                img_input, api_key, style, creativity,
                 neg_prompt, aspect, color_mode, dpi
             ],
             outputs=[prompt_output, quality_report, token_stats],
             api_name="generate"
         ).then(
-            fn=lambda: gr.update(value="Generation complete!"),
+            fn=lambda: "✅ Prompt generated successfully!",
             outputs=status_msg
         )
-        
-        # Modern copy implementation
+
         copy_btn.click(
             fn=copy_text,
             inputs=prompt_output,
             outputs=[prompt_output, status_msg, copy_btn],
-            js="(text) => { if(text) { navigator.clipboard.writeText(text); } return [text]; }"
+            js="""
+            (text) => {
+                if(text) navigator.clipboard.writeText(text);
+                return [text, '✓ Copied!', { variant: 'secondary' }];
+            }
+            """
         )
 
     return app
 
+
 # --- Production Launch ---
 if __name__ == "__main__":
     app = build_interface()