Update app.py

app.py

@@ -1,102 +1,202 @@
 import os
 import gradio as gr
 import google.generativeai as genai
-from PIL import Image
+from PIL import Image, ImageEnhance, ImageFilter
 import io
+import base64
+import json
 import time
+import traceback
+try:
+    import pyperclip
+except ImportError:
+    pyperclip = None
 
-# Hugging Face-safe configuration
-MAX_SIZE = 768
-TIMEOUT = 15
-CONCURRENCY = 2  # Reduced for HF stability
+# --- Environment Configuration ---
+GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
+DEFAULT_PORT = int(os.environ.get("PORT", 7860))
+
+# --- Style Template Optimization ---
+BASE_TEMPLATE = """Describe this design as a concise Flux 1.1 Pro prompt focusing on:
+- Key visual elements
+- Technical specifications
+- Style consistency
+- Functional requirements"""
 
 STYLE_INSTRUCTIONS = {
-    "General": "Describe this design concisely for Flux",
-    "Vector": "Clean vector style with sharp lines",
-    "Realistic": "Photorealistic details",
-    "Kawaii": "Cute cartoon style with pastel colors"
+    "General": BASE_TEMPLATE,
+    "Realistic": f"{BASE_TEMPLATE}\nPHOTOREALISTIC RULES: Use photography terms, texture details, accurate lighting",
+    "Kawaii": f"{BASE_TEMPLATE}\nKAWAII RULES: Rounded shapes, pastel colors, cute expressions",
+    "Vector": f"{BASE_TEMPLATE}\nVECTOR RULES: Clean lines, geometric shapes, B&W gradients",
+    "Silhouette": f"{BASE_TEMPLATE}\nSILHOUETTE RULES: High contrast, minimal details, strong outlines"
+}
+
+# --- Flux Configuration ---
+FLUX_SPECS = {
+    "aspect_ratios": ["1:1", "16:9", "4:3", "9:16"],
+    "formats": ["SVG", "PNG", "PDF"],
+    "color_modes": ["B&W", "CMYK", "RGB"],
+    "dpi_options": [72, 150, 300, 600]
 }
 
-def hf_safe_generate(image, api_key, style):
+# --- Quality Control System ---
+class QualityValidator:
+    VALIDATION_TEMPLATE = """Analyze this Flux prompt:
+1. Score style adherence (1-5)
+2. List technical issues
+3. Suggest improvements
+Respond ONLY as JSON: {"score": x/10, "issues": [], "suggestions": []}"""
+
+    @classmethod
+    def validate(cls, prompt, model):
+        try:
+            response = model.generate_content([cls.VALIDATION_TEMPLATE, prompt])
+            return json.loads(response.text)
+        except:
+            return {"score": 0, "issues": ["Validation failed"], "suggestions": []}
+
+# --- Image Processing Pipeline ---
+def preprocess_image(img):
+    """Convert and enhance uploaded images"""
     try:
-        start_time = time.time()
-        
-        # Convert and resize image
-        img = image.convert("RGB")
-        img.thumbnail((MAX_SIZE, MAX_SIZE))
-        buffer = io.BytesIO()
-        img.save(buffer, format="JPEG", quality=85)
-        
-        # Configure API
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-1.5-pro')
-        
-        # Timeout check
-        if time.time() - start_time > 5:
-            return "Error: Image processing too slow"
+        if isinstance(img, str):  # Handle file paths
+            img = Image.open(img)
+        img = img.convert("RGB")
+        img = ImageEnhance.Contrast(img).enhance(1.2)
+        img = img.filter(ImageFilter.SHARPEN)
+        return img
+    except Exception as e:
+        raise ValueError(f"Image processing error: {str(e)}")
+
+# --- Core Generation Engine ---
+def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi):
+    try:
+        # Validate inputs
+        if not image:
+            raise ValueError("Please upload an image")
             
+        api_key = api_key or GEMINI_KEY
+        if not api_key:
+            raise ValueError("API key required - set in env (GEMINI_KEY) or input field")
+
+        # Initialize model
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel("gemini-1.5-pro")
+
+        # Process image
+        img = preprocess_image(image)
+        img_bytes = io.BytesIO()
+        img.save(img_bytes, format="PNG")
+        img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
+
+        # Build instruction
+        instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
+        instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
+
+        # Generate prompt
         response = model.generate_content(
-            [STYLE_INSTRUCTIONS[style], 
-            {"mime_type": "image/jpeg", "data": buffer.getvalue()},
-            request_options={"timeout": TIMEOUT}
+            contents=[instruction, {"mime_type": "image/png", "data": img_b64}],
+            generation_config={"temperature": creativity}
         )
+        raw_prompt = response.text
+
+        # Quality validation
+        validation = QualityValidator.validate(raw_prompt, model)
+        if validation.get("score", 0) < 7:
+            response = model.generate_content(f"Improve this prompt: {raw_prompt}\nIssues: {validation['issues']}")
+            raw_prompt = response.text
+
+        # Token tracking
+        input_tokens = len(img_b64) // 4  # Approximate base64 token count
+        output_tokens = len(raw_prompt.split())
         
-        return response.text[:500] if response.text else "No response"
-        
+        return {
+            "prompt": raw_prompt,
+            "validation": validation,
+            "stats": {"input": input_tokens, "output": output_tokens}
+        }
+
     except Exception as e:
-        return f"Error: {str(e)}"
-
-# Enhanced interface with original features
-with gr.Blocks(title="Flux Pro HF", theme=gr.themes.Soft()) as app:
-    gr.Markdown("## 🚀 **Flux Prompt Generator**")
-    
-    with gr.Row():
-        with gr.Column():
-            api_key = gr.Textbox(
-                label="🔑 Gemini API Key",
-                type="password",
-                placeholder="Enter your API key..."
-            )
-            style = gr.Dropdown(
-                list(STYLE_INSTRUCTIONS.keys()),
-                value="General",
-                label="🎨 Design Style",
-                interactive=True
-            )
-            image = gr.Image(
-                label="🖼️ Upload Design",
-                type="pil",
-                height=300
-            )
-            
-    with gr.Row():
-        generate_btn = gr.Button("✨ Generate Prompt", variant="primary")
-        copy_btn = gr.Button("📋 Copy")
-        
-    output = gr.Textbox(
-        label="📝 Generated Prompt",
-        lines=5,
-        max_lines=8,
-        interactive=False
-    )
-    
-    # Event handlers
-    generate_btn.click(
-        hf_safe_generate,
-        inputs=[image, api_key, style],
-        outputs=output
-    )
-    
-    copy_btn.click(
-        lambda x: x,
-        inputs=output,
-        outputs=output,
-        _js="(text) => navigator.clipboard.writeText(text)"
-    )
-
-# Fixed queue configuration for Hugging Face
-app.queue(concurrency_count=CONCURRENCY).launch(
-    debug=False,
-    max_threads=1,
-    show_error=True
-)
+        traceback.print_exc()
+        return {"error": str(e)}
+
+# --- UI Components ---
+def create_advanced_controls():
+    with gr.Accordion("⚙️ Advanced Settings", open=False):
+        with gr.Row():
+            creativity = gr.Slider(0.0, 1.0, 0.7, label="Creativity Level")
+            neg_prompt = gr.Textbox(label="🚫 Negative Prompts", placeholder="What to avoid")
+        with gr.Row():
+            aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], label="Aspect Ratio")
+            color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], label="Color Mode")
+            dpi = gr.Dropdown(FLUX_SPECS["dpi_options"], label="Output DPI")
+    return [creativity, neg_prompt, aspect, color_mode, dpi]
+
+# --- Main Interface ---
+def build_interface():
+    with gr.Blocks(title="Flux Pro Generator", theme=gr.themes.Soft()) as app:
+        # Security Section
+        api_key = gr.Textbox(
+            label="🔑 Gemini API Key",
+            value=GEMINI_KEY,
+            type="password",
+            info="Set GEMINI_KEY environment variable for production"
+        )
+
+        # Main Workflow
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                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"
+                )
+                adv_controls = create_advanced_controls()
+                gen_btn = gr.Button("✨ Generate Prompt", variant="primary")
+
+            with gr.Column(scale=2):
+                prompt_output = gr.Textbox(
+                    label="📝 Optimized Prompt",
+                    lines=8,
+                    interactive=False
+                )
+                with gr.Row():
+                    copy_btn = gr.Button("📋 Copy")
+                    history_btn = gr.Button("🕒 History")
+                quality_report = gr.JSON(
+                    label="🔍 Quality Report",
+                    visible=True
+                )
+                token_stats = gr.JSON(
+                    label="🧮 Token Usage",
+                    visible=True
+                )
+
+        # History System
+        history = gr.State([])
+
+        # Event Handling
+        gen_btn.click(
+            lambda *args: generate_prompt(*args),
+            inputs=[img_input, api_key, style] + adv_controls,
+            outputs=[prompt_output, quality_report, token_stats]
+        )
+
+        copy_btn.click(
+            lambda x: pyperclip.copy(x) if pyperclip else None,
+            inputs=prompt_output,
+            outputs=None
+        )
+
+    return app
+
+# --- Production Launch ---
+if __name__ == "__main__":
+    app = build_interface()
+    app.launch()