mroccuper commited on
Commit
70ebe1c
ยท
verified ยท
1 Parent(s): 43c4d6f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +86 -45
app.py CHANGED
@@ -10,7 +10,7 @@ from PIL import Image, ImageEnhance, ImageFilter
10
  # --- Environment Configuration ---
11
  GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
12
  DEFAULT_PORT = int(os.environ.get("PORT", 7860))
13
- API_TIMEOUT = 120 # seconds
14
 
15
  # --- Style Template Optimization ---
16
  BASE_TEMPLATE = """Describe this design as a concise Flux 1.1 Pro prompt focusing on:
@@ -51,36 +51,39 @@ def preprocess_image(img):
51
  # --- Core Generation Engine ---
52
  def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi):
53
  try:
 
54
  if not image:
55
- return {"error": "โŒ Please upload an image"}
56
-
57
  api_key = api_key or GEMINI_KEY
58
  if not api_key:
59
- return {"error": "โŒ API key required - set GEMINI_KEY environment variable or use input field"}
60
 
61
- # Lazy import for Gemini
62
  try:
63
  import google.generativeai as genai
64
  genai.configure(api_key=api_key)
65
  model = genai.GenerativeModel("gemini-1.5-pro")
66
  except ImportError:
67
- return {"error": "โŒ Install Gemini SDK: pip install google-generativeai"}
68
  except Exception as e:
69
  if "authentication" in str(e).lower():
70
- return {"error": "โŒ Invalid API key or authentication error"}
71
- return {"error": f"โŒ API initialization error: {str(e)}"}
 
72
 
73
- # Image processing
 
74
  img = preprocess_image(image)
75
  img_bytes = io.BytesIO()
76
  img.save(img_bytes, format="PNG")
77
  img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
78
 
79
- # Prompt instruction
80
  instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
81
  instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
82
 
83
- # Gemini generation
84
  try:
85
  response = model.generate_content(
86
  contents=[instruction, {"mime_type": "image/png", "data": img_b64}],
@@ -88,41 +91,45 @@ def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color
88
  )
89
  raw_prompt = response.text
90
  except Exception as e:
91
- return {"error": f"โŒ Prompt generation failed: {str(e)}"}
92
 
93
- # Basic validation
94
  validation = {"score": 8, "issues": [], "suggestions": []}
95
- input_tokens = len(img_b64) // 4 # Approximate
 
 
96
  output_tokens = len(raw_prompt.split())
97
-
98
- return raw_prompt, validation, {
99
- "input": input_tokens,
100
- "output": output_tokens
 
101
  }
102
 
103
  except Exception as e:
104
  traceback.print_exc()
105
  return {"error": str(e)}
106
 
107
- # --- Response Formatter ---
108
  def format_generation_response(result):
109
  """Format the response from generate_prompt for the UI"""
110
  if "error" in result:
111
- return result["error"], {}, {}
112
  else:
113
  return result.get("prompt", ""), result.get("validation", {}), result.get("stats", {})
114
 
115
- def update_status(result):
116
- return "โœ… Prompt generated successfully!" if "prompt" in result else result.get("error", "โŒ Unknown error")
 
117
 
118
  # --- Main Interface ---
119
  def build_interface():
120
- with gr.Blocks(title="Flux Pro Generator") as app:
121
  # Header
122
  gr.Markdown("# ๐ŸŽจ Flux Pro Prompt Generator")
123
  gr.Markdown("Generate optimized design prompts from images using Google's Gemini")
124
-
125
- # API Key
126
  api_key = gr.Textbox(
127
  label="๐Ÿ”‘ Gemini API Key",
128
  value=GEMINI_KEY,
@@ -130,42 +137,76 @@ def build_interface():
130
  info="Set GEMINI_KEY environment variable for production"
131
  )
132
 
133
- # Inputs
134
  with gr.Row():
135
  with gr.Column(scale=1):
136
- img_input = gr.Image(label="๐Ÿ–ผ๏ธ Upload Design", type="pil", sources=["upload"], interactive=True)
137
- style = gr.Dropdown(list(STYLE_INSTRUCTIONS.keys()), value="General", label="๐ŸŽจ Target Style")
 
 
 
 
 
 
 
 
 
 
 
138
  with gr.Accordion("โš™๏ธ Advanced Settings", open=False):
139
- creativity = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Creativity Level")
140
  neg_prompt = gr.Textbox(label="๐Ÿšซ Negative Prompts", placeholder="What to avoid")
141
  aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
142
  color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], value="RGB", label="Color Mode")
143
  dpi = gr.Dropdown([str(d) for d in FLUX_SPECS["dpi_options"]], value="300", label="Output DPI")
 
144
  gen_btn = gr.Button("โœจ Generate Prompt", variant="primary")
145
 
146
  with gr.Column(scale=2):
147
- prompt_output = gr.Textbox(label="๐Ÿ“ Optimized Prompt", lines=8, interactive=True, show_copy_button=True)
 
 
 
 
 
148
  status_msg = gr.Textbox(label="Status", visible=True)
149
- quality_report = gr.JSON(label="๐Ÿ” Quality Report", visible=True)
150
- token_stats = gr.JSON(label="๐Ÿงฎ Token Usage", visible=True)
151
-
152
- # Event bindings
 
 
 
 
 
 
 
 
 
153
  gen_btn.click(
154
  fn=generate_prompt,
155
- inputs=[img_input, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi],
156
- outputs=None
 
 
 
 
157
  ).then(
158
- fn=format_generation_response,
159
- inputs=None,
160
- outputs=[prompt_output, quality_report, token_stats]
161
- ).then(
162
- fn=update_status,
163
- inputs=[prompt_output],
164
- outputs=[status_msg]
 
 
 
165
  )
 
166
  return app
167
 
168
- # --- Launch App ---
169
  if __name__ == "__main__":
170
  app = build_interface()
171
- app.launch(server_port=DEFAULT_PORT)
 
10
  # --- Environment Configuration ---
11
  GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
12
  DEFAULT_PORT = int(os.environ.get("PORT", 7860))
13
+ API_TIMEOUT = 30 # seconds
14
 
15
  # --- Style Template Optimization ---
16
  BASE_TEMPLATE = """Describe this design as a concise Flux 1.1 Pro prompt focusing on:
 
51
  # --- Core Generation Engine ---
52
  def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi):
53
  try:
54
+ # Validate inputs
55
  if not image:
56
+ return {"error": "Please upload an image"}
57
+
58
  api_key = api_key or GEMINI_KEY
59
  if not api_key:
60
+ return {"error": "API key required - set in env (GEMINI_KEY) or input field"}
61
 
62
+ # Import and configure Gemini only when needed
63
  try:
64
  import google.generativeai as genai
65
  genai.configure(api_key=api_key)
66
  model = genai.GenerativeModel("gemini-1.5-pro")
67
  except ImportError:
68
+ return {"error": "Failed to import google.generativeai. Install with: pip install google-generativeai"}
69
  except Exception as e:
70
  if "authentication" in str(e).lower():
71
+ return {"error": "Invalid API key or authentication error"}
72
+ else:
73
+ return {"error": f"API initialization error: {str(e)}"}
74
 
75
+ # Process image with timeout protection
76
+ start_time = time.time()
77
  img = preprocess_image(image)
78
  img_bytes = io.BytesIO()
79
  img.save(img_bytes, format="PNG")
80
  img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
81
 
82
+ # Build instruction
83
  instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
84
  instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
85
 
86
+ # Generate prompt with timeout protection
87
  try:
88
  response = model.generate_content(
89
  contents=[instruction, {"mime_type": "image/png", "data": img_b64}],
 
91
  )
92
  raw_prompt = response.text
93
  except Exception as e:
94
+ return {"error": f"Generation failed: {str(e)}"}
95
 
96
+ # Simple quality validation
97
  validation = {"score": 8, "issues": [], "suggestions": []}
98
+
99
+ # Token tracking
100
+ input_tokens = len(img_b64) // 4 # Approximate base64 token count
101
  output_tokens = len(raw_prompt.split())
102
+
103
+ return {
104
+ "prompt": raw_prompt,
105
+ "validation": validation,
106
+ "stats": {"input": input_tokens, "output": output_tokens}
107
  }
108
 
109
  except Exception as e:
110
  traceback.print_exc()
111
  return {"error": str(e)}
112
 
113
+ # --- UI Response Formatting ---
114
  def format_generation_response(result):
115
  """Format the response from generate_prompt for the UI"""
116
  if "error" in result:
117
+ return result["error"], None, None
118
  else:
119
  return result.get("prompt", ""), result.get("validation", {}), result.get("stats", {})
120
 
121
+ # Modern copy function using Gradio's JavaScript API
122
+ def copy_text(text):
123
+ return gr.update(value=text), f"โœ“ Copied: '{text[:20]}...'", gr.Button.update(variant="secondary")
124
 
125
  # --- Main Interface ---
126
  def build_interface():
127
+ with gr.Blocks(title="Flux Pro Generator", theme="soft") as app:
128
  # Header
129
  gr.Markdown("# ๐ŸŽจ Flux Pro Prompt Generator")
130
  gr.Markdown("Generate optimized design prompts from images using Google's Gemini")
131
+
132
+ # Security Section
133
  api_key = gr.Textbox(
134
  label="๐Ÿ”‘ Gemini API Key",
135
  value=GEMINI_KEY,
 
137
  info="Set GEMINI_KEY environment variable for production"
138
  )
139
 
140
+ # Main Workflow
141
  with gr.Row():
142
  with gr.Column(scale=1):
143
+ img_input = gr.Image(
144
+ label="๐Ÿ–ผ๏ธ Upload Design",
145
+ type="pil",
146
+ sources=["upload"],
147
+ interactive=True
148
+ )
149
+ style = gr.Dropdown(
150
+ list(STYLE_INSTRUCTIONS.keys()),
151
+ value="General",
152
+ label="๐ŸŽจ Target Style"
153
+ )
154
+
155
+ # Advanced Settings
156
  with gr.Accordion("โš™๏ธ Advanced Settings", open=False):
157
+ creativity = gr.Slider(0.0, 1.0, 0.7, label="Creativity Level")
158
  neg_prompt = gr.Textbox(label="๐Ÿšซ Negative Prompts", placeholder="What to avoid")
159
  aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
160
  color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], value="RGB", label="Color Mode")
161
  dpi = gr.Dropdown([str(d) for d in FLUX_SPECS["dpi_options"]], value="300", label="Output DPI")
162
+
163
  gen_btn = gr.Button("โœจ Generate Prompt", variant="primary")
164
 
165
  with gr.Column(scale=2):
166
+ prompt_output = gr.Textbox(
167
+ label="๐Ÿ“ Optimized Prompt",
168
+ lines=8,
169
+ interactive=True,
170
+ show_copy_button=True # Modern Gradio has built-in copy button
171
+ )
172
  status_msg = gr.Textbox(label="Status", visible=True)
173
+ with gr.Row():
174
+ copy_btn = gr.Button("๐Ÿ“‹ Copy to Clipboard", variant="secondary")
175
+
176
+ quality_report = gr.JSON(
177
+ label="๐Ÿ” Quality Report",
178
+ visible=True
179
+ )
180
+ token_stats = gr.JSON(
181
+ label="๐Ÿงฎ Token Usage",
182
+ visible=True
183
+ )
184
+
185
+ # Event Handling
186
  gen_btn.click(
187
  fn=generate_prompt,
188
+ inputs=[
189
+ img_input, api_key, style, creativity,
190
+ neg_prompt, aspect, color_mode, dpi
191
+ ],
192
+ outputs=[prompt_output, quality_report, token_stats],
193
+ api_name="generate"
194
  ).then(
195
+ fn=lambda: gr.update(value="Generation complete!"),
196
+ outputs=status_msg
197
+ )
198
+
199
+ # Modern copy implementation
200
+ copy_btn.click(
201
+ fn=copy_text,
202
+ inputs=prompt_output,
203
+ outputs=[prompt_output, status_msg, copy_btn],
204
+ js="(text) => { if(text) { navigator.clipboard.writeText(text); } return [text]; }"
205
  )
206
+
207
  return app
208
 
209
+ # --- Production Launch ---
210
  if __name__ == "__main__":
211
  app = build_interface()
212
+ app.launch()