Update app.py

app.py

@@ -1,210 +1,76 @@
 import os
 import gradio as gr
 import google.generativeai as genai
-from PIL import Image, ImageEnhance, ImageFilter
+from PIL import Image
 import io
 import base64
-import json
 import time
-import traceback
-try:
-    import pyperclip
-except ImportError:
-    pyperclip = None
 
-# --- 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"""
+# Configuration optimized for Hugging Face
+MAX_RETRIES = 3
+TIMEOUT = 10  # Keep under 20s Hugging Face limit
 
+# Simplified style instructions for faster processing
 STYLE_INSTRUCTIONS = {
-    "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]
+    "General": "Describe as Flux prompt focusing on key elements",
+    "Vector": "Clean vector style with sharp lines and B&W",
+    "Realistic": "Photorealistic details with accurate textures",
 }
 
-# --- 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": []}"""
+def optimize_image(img):
+    """Fast image optimization for web"""
+    img = img.convert("RGB")
+    img.thumbnail((512, 512))  # Reduce image size
+    return img
 
-    @classmethod
-    def validate(cls, prompt, model):
+def generate_prompt(image, api_key, style):
+    for attempt in range(MAX_RETRIES):
         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:
-        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")
+            start_time = time.time()
             
-        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(
-            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 {
-            "prompt": raw_prompt,
-            "validation": validation,
-            "stats": {"input": input_tokens, "output": output_tokens}
-        }
-
-    except Exception as e:
-        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(
-        server_name="0.0.0.0",
-        server_port=DEFAULT_PORT,
-        share=False,
-        auth=(
-            os.environ.get("APP_USER", "admin"),
-            os.environ.get("APP_PWD", "admin")
-        ) if os.environ.get("ENV") == "prod" else None
-    )
+            # Configure API
+            genai.configure(api_key=api_key)
+            model = genai.GenerativeModel('gemini-pro-vision')
+            
+            # Optimize image
+            img = optimize_image(image)
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='JPEG', quality=85)
+            
+            # Timeout control
+            if time.time() - start_time > TIMEOUT - 2:
+                raise TimeoutError("Image processing too slow")
+                
+            response = model.generate_content([
+                STYLE_INSTRUCTIONS[style],
+                {"mime_type": "image/jpeg", "data": img_byte_arr.getvalue()}
+            ], request_options={"timeout": TIMEOUT})
+            
+            return response.text
+            
+        except Exception as e:
+            if attempt == MAX_RETRIES - 1:
+                return f"Error: {str(e)}"
+            time.sleep(1)
+
+# Hugging Face optimized interface
+with gr.Blocks() as app:
+    gr.Markdown("## 🚀 HF-Friendly Flux Generator")
+    
+    with gr.Row():
+        api_key = gr.Textbox(label="Gemini Key", type="password")
+        style = gr.Dropdown(list(STYLE_INSTRUCTIONS), label="Style")
+    
+    image_input = gr.Image(type="pil", label="Upload Design")
+    generate_btn = gr.Button("Generate", variant="primary")
+    output = gr.Textbox(label="Prompt", interactive=False)
+
+    generate_btn.click(
+        generate_prompt,
+        inputs=[image_input, api_key, style],
+        outputs=output,
+        concurrency_limit=3  # Better HF resource management
+    )
+
+# HF-specific launch settings
+app.queue(concurrency_count=3, max_size=20).launch(debug=True)