Update app.py

app.py

@@ -3,6 +3,7 @@ import google.generativeai as genai
 from PIL import Image
 import io
 import base64
+import requests
 
 # List of popular styles
 STYLES = [
@@ -27,40 +28,48 @@ def enhance_prompt(prompt, style):
     
     return response.text
 
-def generate_image(enhanced_prompt, style, negative_prompt):
-    model = genai.GenerativeModel("gemini-2.0-flash")
-    full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
+def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
+    url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
     
-    response = model.generate_content(full_prompt)
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {stability_api_key}"
+    }
     
-    if response.parts:
-        for part in response.parts:
-            if part.image:
-                image_bytes = part.image.to_bytes()
-                image = Image.open(io.BytesIO(image_bytes))
-                
-                # Convert PIL Image to base64 string
-                buffered = io.BytesIO()
-                image.save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode()
-                
-                return f"data:image/png;base64,{img_str}"
+    payload = {
+        "model": "sd3.5-large-turbo",
+        "prompt": f"{enhanced_prompt}, Style: {style}",
+        "negative_prompt": negative_prompt,
+        "width": 1024,
+        "height": 1024,
+        "num_images": 1,
+        "steps": 30,
+        "cfg_scale": 7.5
+    }
     
-    return "Image generation failed"
+    response = requests.post(url, json=payload, headers=headers)
+    
+    if response.status_code == 200:
+        image_data = response.json()['images'][0]['image_base64']
+        return f"data:image/png;base64,{image_data}"
+    else:
+        return f"Image generation failed: {response.text}"
 
-def process_and_generate(api_key, prompt, style, negative_prompt):
-    genai.configure(api_key=api_key)
+def process_and_generate(google_api_key, stability_api_key, prompt, style, negative_prompt):
+    genai.configure(api_key=google_api_key)
     
     enhanced_prompt = enhance_prompt(prompt, style)
-    image_url = generate_image(enhanced_prompt, style, negative_prompt)
+    image_url = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
     return image_url, enhanced_prompt
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Google Imagen 3 Image Generator")
+    gr.Markdown("# Stability AI Image Generator with Google Gemini Prompt Enhancement")
     
     with gr.Row():
         with gr.Column(scale=1):
-            api_key = gr.Textbox(label="Google AI API Key", type="password")
+            google_api_key = gr.Textbox(label="Google AI API Key", type="password")
+            stability_api_key = gr.Textbox(label="Stability AI API Key", type="password")
             prompt = gr.Textbox(label="Prompt")
             style = gr.Dropdown(label="Style", choices=STYLES)
             negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
@@ -72,7 +81,7 @@ with gr.Blocks() as demo:
     
     generate_btn.click(
         process_and_generate,
-        inputs=[api_key, prompt, style, negative_prompt],
+        inputs=[google_api_key, stability_api_key, prompt, style, negative_prompt],
         outputs=[image_output, enhanced_prompt_output]
     )