mroccuper commited on
Commit
ca91c2d
ยท
verified ยท
1 Parent(s): cce6f28

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -77
app.py CHANGED
@@ -51,39 +51,36 @@ 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
- # 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,15 +88,13 @@ def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color
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,
@@ -110,26 +105,25 @@ def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color
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,76 +131,40 @@ def build_interface():
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()
 
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
  )
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 {
99
  "prompt": raw_prompt,
100
  "validation": validation,
 
105
  traceback.print_exc()
106
  return {"error": str(e)}
107
 
108
+ # --- Response Formatter ---
109
  def format_generation_response(result):
110
  """Format the response from generate_prompt for the UI"""
111
  if "error" in result:
112
+ return result["error"], {}, {}
113
  else:
114
  return result.get("prompt", ""), result.get("validation", {}), result.get("stats", {})
115
 
116
+ def update_status(result):
117
+ return "โœ… Prompt generated successfully!" if "prompt" in result else result.get("error", "โŒ Unknown error")
 
118
 
119
  # --- Main Interface ---
120
  def build_interface():
121
+ with gr.Blocks(title="Flux Pro Generator") as app:
122
  # Header
123
  gr.Markdown("# ๐ŸŽจ Flux Pro Prompt Generator")
124
  gr.Markdown("Generate optimized design prompts from images using Google's Gemini")
125
+
126
+ # API Key
127
  api_key = gr.Textbox(
128
  label="๐Ÿ”‘ Gemini API Key",
129
  value=GEMINI_KEY,
 
131
  info="Set GEMINI_KEY environment variable for production"
132
  )
133
 
134
+ # Inputs
135
  with gr.Row():
136
  with gr.Column(scale=1):
137
+ img_input = gr.Image(label="๐Ÿ–ผ๏ธ Upload Design", type="pil", sources=["upload"], interactive=True)
138
+ style = gr.Dropdown(list(STYLE_INSTRUCTIONS.keys()), value="General", label="๐ŸŽจ Target Style")
 
 
 
 
 
 
 
 
 
 
 
139
  with gr.Accordion("โš™๏ธ Advanced Settings", open=False):
140
+ creativity = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Creativity Level")
141
  neg_prompt = gr.Textbox(label="๐Ÿšซ Negative Prompts", placeholder="What to avoid")
142
  aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
143
  color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], value="RGB", label="Color Mode")
144
  dpi = gr.Dropdown([str(d) for d in FLUX_SPECS["dpi_options"]], value="300", label="Output DPI")
 
145
  gen_btn = gr.Button("โœจ Generate Prompt", variant="primary")
146
 
147
  with gr.Column(scale=2):
148
+ prompt_output = gr.Textbox(label="๐Ÿ“ Optimized Prompt", lines=8, interactive=True, show_copy_button=True)
 
 
 
 
 
149
  status_msg = gr.Textbox(label="Status", visible=True)
150
+ quality_report = gr.JSON(label="๐Ÿ” Quality Report", visible=True)
151
+ token_stats = gr.JSON(label="๐Ÿงฎ Token Usage", visible=True)
152
+
153
+ # Event bindings
 
 
 
 
 
 
 
 
 
154
  gen_btn.click(
155
  fn=generate_prompt,
156
+ inputs=[img_input, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi],
 
 
 
157
  outputs=[prompt_output, quality_report, token_stats],
158
  api_name="generate"
159
  ).then(
160
+ fn=lambda r: update_status({"prompt": r}) if isinstance(r, str) else update_status(r),
161
+ inputs=[prompt_output],
162
  outputs=status_msg
163
  )
 
 
 
 
 
 
 
 
164
 
165
  return app
166
 
167
+ # --- Launch App ---
168
  if __name__ == "__main__":
169
  app = build_interface()
170
+ app.launch(server_port=DEFAULT_PORT)