Update app.py

app.py

@@ -6,166 +6,205 @@ import io
 import base64
 import json
 import time
+import traceback
 try:
     import pyperclip
 except ImportError:
-    os.system('pip install pyperclip')
-    import pyperclip
+    pyperclip = None
 
-# --- Security First Configuration ---
-GEMINI_KEY = os.environ.get("GEMINI_KEY", "")  # Load from environment variable
+# --- Environment Configuration ---
+GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
+DEFAULT_PORT = int(os.environ.get("PORT", 7860))
 
-# --- Template Optimization System ---
-BASE_TEMPLATE = (
-    "Describe this design as a single, cohesive, and concise prompt suitable for Flux 1.1 Pro. "
-    "Focus on key elements such as text, symbols, layout, and overall style."
-)
+# --- 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": BASE_TEMPLATE,
-    "Realistic": (
-        f"{BASE_TEMPLATE} Focus on photorealistic details: textures, lighting, depth. "
-        "Avoid illustrations/cartoon styles. Use technical photography terms."
-    ),
-    "Kawaii/Cartoon": (
-        f"{BASE_TEMPLATE} Emphasize cute, rounded shapes, playful expressions, "
-        "vibrant colors. Use anime/kawaii terminology."
-    ),
-    "Vector": (
-        f"{BASE_TEMPLATE} Specify clean vector style with sharp lines, geometric shapes. "
-        "Black/white only with gray gradients. Use design technical terms."
-    ),
-    "Silhouette": (
-        f"{BASE_TEMPLATE} Use high-contrast black/white silhouettes. "
-        "Focus on bold shapes and outlines. Eliminate interior details."
-    ),
+    "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 Engine ---
+# --- Flux Configuration ---
 FLUX_SPECS = {
-    "aspect_ratios": ["1:1", "16:9", "4:3", "9:16", "Custom"],
-    "output_formats": ["SVG", "PNG-300dpi", "PDF", "EPS"],
-    "color_modes": ["B&W", "CMYK", "Spot Colors", "RGB"],
-    "complexity_levels": ["Minimal", "Medium", "High", "Ultra"]
+    "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]
 }
 
-# --- Quality Control Systems ---
+# --- Quality Control System ---
 class QualityValidator:
-    VALIDATION_PROMPT = """Analyze this Flux prompt on:
-1. Technical specificity (1-5)
-2. Style adherence (1-5)
-3. Flux compatibility (1-5)
-Respond ONLY as JSON: {"total": x/15, "issues": [list]}"""
+    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_PROMPT, prompt])
+            response = model.generate_content([cls.VALIDATION_TEMPLATE, prompt])
             return json.loads(response.text)
         except:
-            return {"total": 0, "issues": ["Validation failed"]}
+            return {"score": 0, "issues": ["Validation failed"], "suggestions": []}
 
 # --- Image Processing Pipeline ---
 def preprocess_image(img):
-    """Enhance image quality before analysis"""
-    img = img.convert("RGB")
-    img = ImageEnhance.Contrast(img).enhance(1.2)
-    img = img.filter(ImageFilter.SHARPEN)
-    return img
-
-# --- Core Application Logic ---
-def generate_prompt(image, api_key, style, creativity, neg_prompt, flux_specs):
-    genai.configure(api_key=api_key or GEMINI_KEY)
-    model = genai.GenerativeModel("gemini-1.5-pro", generation_config={
-        "temperature": creativity,
-        "top_p": 0.95
-    })
-
-    # Pre-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}\nFLUX SPECS: {flux_specs}"
-
-    # Generate initial prompt
-    response = model.generate_content([instruction, {"mime_type": "image/png", "data": img_b64}])
-    raw_prompt = response.text
-
-    # Quality validation
-    validation = QualityValidator.validate(raw_prompt, model)
-    if validation["total"] < 10:  # Regenerate if low quality
-        response = model.generate_content([
-            f"Improve this prompt addressing: {validation['issues']}\n\n{raw_prompt}"
-        ])
+    """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")
+            
+        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
 
-    return raw_prompt, validation
+        # 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_params():
-    with gr.Accordion("⚙️ Advanced Parameters", open=False):
+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")
-            neg_prompt = gr.Textbox(label="🚫 Avoid in Prompt", placeholder="e.g., no text, avoid gradients")
+            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"], value="1:1", label="Aspect Ratio")
-            fmt = gr.Dropdown(FLUX_SPECS["output_formats"], value="SVG", label="Output Format")
-            color = gr.Dropdown(FLUX_SPECS["color_modes"], value="B&W", label="Color Mode")
-            complexity = gr.Dropdown(FLUX_SPECS["complexity_levels"], value="Medium", label="Complexity")
-    return [creativity, neg_prompt, aspect, fmt, color, complexity]
+            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]
 
-# --- Gradio Interface ---
+# --- Main Interface ---
 def build_interface():
-    with gr.Blocks(title="Flux Pro Generator") as app:
-        # Security First
-        api_key = gr.Textbox(label="🔑 Gemini API Key", value=GEMINI_KEY, type="password", 
-                           info="Set GEMINI_KEY environment variable for production")
+    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():
+        with gr.Row(variant="panel"):
             with gr.Column(scale=1):
-                img_input = gr.Image(label="🖼️ Upload Design", type="pil", sources=["upload", "clipboard"])
-                style = gr.Dropdown(list(STYLE_INSTRUCTIONS), value="General", label="🎨 Target Style")
-                adv_params = create_advanced_params()
-                
+                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)
+                prompt_output = gr.Textbox(
+                    label="📝 Optimized Prompt",
+                    lines=8,
+                    interactive=False
+                )
                 with gr.Row():
-                    gen_btn = gr.Button("✨ Generate", variant="primary")
                     copy_btn = gr.Button("📋 Copy")
-                quality_report = gr.JSON(label="🔍 Quality Report")
-
-        # Enterprise Features
-        token_counter = gr.Textbox(label="🧮 Token Usage", interactive=False)
+                    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(
-            generate_prompt,
-            inputs=[img_input, api_key, style] + adv_params,
-            outputs=[prompt_output, quality_report]
+            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), x),
+            lambda x: pyperclip.copy(x) if pyperclip else None,
             inputs=prompt_output,
-            outputs=prompt_output
+            outputs=None
         )
 
     return app
 
-# --- Production Deployment ---
+# --- Production Launch ---
 if __name__ == "__main__":
     app = build_interface()
     app.launch(
         server_name="0.0.0.0",
-        server_port=int(os.getenv("PORT", 7860)),
+        server_port=DEFAULT_PORT,
         share=False,
         auth=(
-            os.getenv("APP_USER", "admin"), 
-            os.getenv("APP_PWD", "admin")
-        ) if os.getenv("ENV") == "prod" else None
+            os.environ.get("APP_USER", "admin"),
+            os.environ.get("APP_PWD", "admin")
+        ) if os.environ.get("ENV") == "prod" else None
     )