Spaces:
Running
Running
| import os | |
| import gradio as gr | |
| import traceback | |
| from PIL import Image, ImageEnhance, ImageFilter | |
| import io | |
| import base64 | |
| import google.generativeai as genai | |
| # --- Environment Configuration --- | |
| GEMINI_KEY = os.environ.get("GEMINI_KEY", "") | |
| # --- 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}\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"], | |
| "color_modes": ["B&W", "CMYK", "RGB"], | |
| "dpi_options": [72, 150, 300, 600] | |
| } | |
| # --- Image Processing Pipeline --- | |
| def preprocess_image(img): | |
| 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 failed: {str(e)}") | |
| # --- Core Generation Engine --- | |
| def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi): | |
| try: | |
| # Step 1: Input Validation | |
| if not image: | |
| return "", "⚠️ Please upload an image." | |
| api_key = api_key or GEMINI_KEY | |
| if not api_key: | |
| return "", "🔑 API key required - set in env (GEMINI_KEY) or input field." | |
| # Step 2: Gemini Setup | |
| try: | |
| genai.configure(api_key=api_key) | |
| model = genai.GenerativeModel("gemini-1.5-pro") | |
| except ImportError: | |
| return "", "🚫 Failed to import google.generativeai. Install with: pip install google-generativeai" | |
| except Exception as e: | |
| if "authentication" in str(e).lower(): | |
| return "", "🔐 Invalid API key or authentication error." | |
| else: | |
| return "", f"⚙️ API initialization error: {str(e)}" | |
| # 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 "", f"🖼️ Image preparation failed: {str(e)}" | |
| # 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 "", "🛠️ Invalid style selected. Please choose from available options." | |
| # Step 5: Call Gemini API | |
| try: | |
| response = model.generate_content( | |
| contents=[instruction, {"mime_type": "image/png", "data": img_b64}], | |
| generation_config={"temperature": creativity} | |
| ) | |
| raw_prompt = response.text | |
| except Exception as e: | |
| return "", f"🤖 Prompt generation failed: {str(e)}" | |
| return raw_prompt, "✅ Prompt generated successfully!" | |
| except Exception as e: | |
| traceback.print_exc() | |
| return "", f"💥 Unexpected error: {str(e)}" | |
| # --- Main Interface --- | |
| def build_interface(): | |
| global STYLE_INSTRUCTIONS # Ensure access to global dict | |
| with gr.Blocks(title="Flux Pro Generator") as app: | |
| gr.Markdown("# 🎨 Flux Pro Prompt Generator") | |
| gr.Markdown("Generate optimized design prompts from images using Google's Gemini.") | |
| with gr.Row(): | |
| # Left Panel: Image Upload & API Key | |
| 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 | |
| ) | |
| # Right Panel: Prompt Output + Controls Below | |
| with gr.Column(scale=2): | |
| status_msg = gr.Textbox(label="📢 Status", interactive=False) | |
| prompt_output = gr.Textbox( | |
| label="📝 Optimized Prompt", | |
| lines=8, | |
| interactive=True | |
| ) | |
| # Controls under prompt box | |
| style = gr.Dropdown( | |
| list(STYLE_INSTRUCTIONS.keys()), | |
| value="General", | |
| label="🎨 Target Style" | |
| ) | |
| with gr.Accordion("⚙️ Advanced Settings", open=False): | |
| 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") | |
| # Event Handling | |
| gen_btn.click( | |
| fn=generate_prompt, | |
| inputs=[ | |
| img_input, api_key, style, creativity, | |
| neg_prompt, aspect, color_mode, dpi | |
| ], | |
| outputs=[prompt_output, status_msg], | |
| api_name="generate" | |
| ) | |
| return app | |
| # --- Production Launch --- | |
| if __name__ == "__main__": | |
| app = build_interface() | |
| app.launch() |