Spaces:
Running
Running
Update app.py
Browse files
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 =
|
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": "
|
56 |
-
|
57 |
api_key = api_key or GEMINI_KEY
|
58 |
if not api_key:
|
59 |
-
return {"error": "
|
60 |
|
61 |
-
#
|
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": "
|
68 |
except Exception as e:
|
69 |
if "authentication" in str(e).lower():
|
70 |
-
return {"error": "
|
71 |
-
|
|
|
72 |
|
73 |
-
#
|
|
|
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 |
-
#
|
80 |
instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
|
81 |
instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
|
82 |
|
83 |
-
#
|
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"
|
92 |
|
93 |
-
#
|
94 |
validation = {"score": 8, "issues": [], "suggestions": []}
|
95 |
-
|
|
|
|
|
96 |
output_tokens = len(raw_prompt.split())
|
97 |
-
|
98 |
-
return
|
99 |
-
"
|
100 |
-
"
|
|
|
101 |
}
|
102 |
|
103 |
except Exception as e:
|
104 |
traceback.print_exc()
|
105 |
return {"error": str(e)}
|
106 |
|
107 |
-
# --- Response
|
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 |
-
|
116 |
-
|
|
|
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 |
-
#
|
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 |
-
#
|
134 |
with gr.Row():
|
135 |
with gr.Column(scale=1):
|
136 |
-
img_input = gr.Image(
|
137 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
138 |
with gr.Accordion("โ๏ธ Advanced Settings", open=False):
|
139 |
-
creativity = gr.Slider(0.0, 1.0,
|
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(
|
|
|
|
|
|
|
|
|
|
|
148 |
status_msg = gr.Textbox(label="Status", visible=True)
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
153 |
gen_btn.click(
|
154 |
fn=generate_prompt,
|
155 |
-
inputs=[
|
156 |
-
|
|
|
|
|
|
|
|
|
157 |
).then(
|
158 |
-
fn=
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
|
|
|
|
|
|
165 |
)
|
|
|
166 |
return app
|
167 |
|
168 |
-
# --- Launch
|
169 |
if __name__ == "__main__":
|
170 |
app = build_interface()
|
171 |
-
app.launch(
|
|
|
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()
|