Update app.py

app.py

@@ -1,14 +1,13 @@
 import os
 import base64
-import zipfile
-import tempfile
-import uuid
 import json
 import requests
-from datetime import datetime
-from PIL import Image
 import gradio as gr
 from io import BytesIO
+from PIL import Image
+import uuid
+from datetime import datetime
+import zipfile
 
 # Persistent storage setup
 BASE_DIR = os.path.abspath(os.getcwd())
@@ -49,7 +48,7 @@ def get_history(username=None):
         metadata = [m for m in metadata if m["username"] == username]
     else:
         return [], []
-    return [(m["path"], f"{m['username']} - {m['timestamp']}") for m in metadata], [(m["path"], f"{m['username']} - {m['timestamp']}") for m in metadata]
+    return [(m["path"], f"{m['username']} - {m['timestamp']}") for m in metadata]
 
 def download_history():
     with zipfile.ZipFile(ZIP_PATH, "w") as zipf:
@@ -60,11 +59,14 @@ def download_history():
         zipf.write(HISTORY_METADATA, arcname="metadata.json")
     return ZIP_PATH
 
-def generate_image(prompt, bytez_api_key, aspect_ratio="1:1"):
+def generate_image(prompt, bytez_api_key, username):
+    if not username.strip():
+        raise gr.Error("Username cannot be empty")
+    
     api_key = bytez_api_key.strip() if bytez_api_key and bytez_api_key.strip() != "" else os.environ.get("BYTEZ_API_KEY")
     
     if not api_key:
-        raise ValueError("No API key provided. Set BYTEZ_API_KEY environment variable or enter in UI.")
+        raise gr.Error("No API key provided. Set BYTEZ_API_KEY environment variable or enter in UI.")
     
     headers = {
         "Authorization": api_key,
@@ -73,7 +75,7 @@ def generate_image(prompt, bytez_api_key, aspect_ratio="1:1"):
     
     payload = {
         "text": prompt,
-        "aspect_ratio": aspect_ratio
+        #"aspect_ratio": "1:1"  # Fixed to 1:1 ratio as requested
     }
 
     try:
@@ -87,55 +89,40 @@ def generate_image(prompt, bytez_api_key, aspect_ratio="1:1"):
         
         result = response.json()
         if not result.get("results") or not isinstance(result["results"], list) or len(result["results"]) == 0:
-            raise ValueError("No results returned from API")
+            raise gr.Error("No results returned from API")
         
         # Decode base64 image
         img_data = base64.b64decode(result["results"][0]["base64"])
         img = Image.open(BytesIO(img_data))
-        return img
+        
+        # Save to history
+        saved_path = save_image(img, username, prompt)
+        
+        return saved_path
     
     except requests.exceptions.RequestException as e:
         error_msg = f"API request failed: {str(e)}"
         if hasattr(e, 'response') and e.response is not None:
             error_msg += f"\nStatus: {e.response.status_code}\nResponse: {e.response.text}"
-        raise Exception(error_msg)
+        raise gr.Error(error_msg)
     except Exception as e:
-        raise Exception(f"Image generation failed: {str(e)}")
+        raise gr.Error(f"Image generation failed: {str(e)}")
 
-def process_prompt(prompt, bytez_api_key, username, aspect_ratio):
-    if not username.strip():
-        raise ValueError("Username cannot be empty")
-    
-    # Generate image using Imagen API
-    result_img = generate_image(prompt, bytez_api_key, aspect_ratio)
-    
-    # Save to history
-    saved_path = save_image(result_img, username, prompt)
-    
-    # Return for gallery display
-    return Image.open(saved_path), [(saved_path, f"{username}: {prompt}")]
-
-with gr.Blocks() as demo:
+with gr.Blocks(title="Imagen 4.0 Generator") as demo:
     gr.HTML("""
-    <div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
-        <div style="background-color: var(--block-background-fill); border-radius: 8px">
-            <img src="https://api.bytez.com/favicon.ico" style="width: 100px; height: 100px;">
+    <div style='display: flex; align-items: center; justify-content: center; gap: 20px; margin-bottom: 20px'>
+        <div style="background-color: var(--block-background-fill); border-radius: 8px; padding: 10px">
+            <img src="https://api.bytez.com/favicon.ico" style="width: 80px; height: 80px;">
         </div>
         <div>
-            <h1>Imagen 4.0 Fast Generator</h1>
-            <p>Text-to-image generation using Google's Imagen 4.0 model</p>
+            <h1 style="margin-bottom: 5px">Imagen 4.0 Fast Generator</h1>
+            <p style="margin-top: 0; color: var(--body-text-color-subdued)">Text-to-image generation using Google's Imagen 4.0 model</p>
         </div>
     </div>
     """)
 
     with gr.Row():
-        with gr.Column():
-            bytez_api_key = gr.Textbox(
-                lines=1, 
-                placeholder="Enter Bytez API Key (optional)", 
-                label="Bytez API Key",
-                type="password"
-            )
+        with gr.Column(scale=1):
             username_input = gr.Textbox(
                 lines=1, 
                 placeholder="Enter your username", 
@@ -143,62 +130,56 @@ with gr.Blocks() as demo:
                 info="Required for saving history"
             )
             prompt_input = gr.Textbox(
-                lines=3, 
+                lines=4, 
                 placeholder="A Raspberry Pi SBC on a wooden desk with electronic components", 
                 label="Prompt",
                 info="Be descriptive for best results"
             )
-            aspect_ratio = gr.Dropdown(
-                choices=["1:1", "16:9", "9:16"],
-                value="1:1",
-                label="Aspect Ratio"
+            submit_btn = gr.Button("✨ Generate Image", variant="primary", size="lg")
+            
+            gr.Examples(
+                examples=[
+                    ["A cyberpunk cat wearing neon goggles, digital art style"],
+                    ["Mountain landscape at sunset with northern lights, photorealistic"],
+                    ["Steampunk airship floating above Victorian city, detailed illustration"],
+                    ["Minimalist logo of a bird made from geometric shapes, blue and white"]
+                ],
+                inputs=prompt_input,
+                label="💡 Example Prompts"
             )
-            submit_btn = gr.Button("✨ Generate Image", variant="primary")
 
-        with gr.Column():
-            output_image = gr.Image(
-                label="Generated Image",
-                interactive=False
-            )
-            history_gallery = gr.Gallery(
-                label="Your History",
-                columns=2,
-                preview=True
-            )
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Generated Image", height=512)
+            history_gallery = gr.Gallery(label="Your History", columns=3, preview=True)
+            
             with gr.Row():
                 download_zip_btn = gr.Button("📥 Download History ZIP")
-                download_zip_file = gr.File(label="Download ZIP")
+                download_zip_file = gr.File(label="Download ZIP", visible=False)
 
     # Event handlers
     submit_btn.click(
-        fn=process_prompt,
-        inputs=[prompt_input, bytez_api_key, username_input, aspect_ratio],
-        outputs=[output_image, history_gallery],
-        api_name="generate"
+        fn=generate_image,
+        inputs=[prompt_input, bytez_api_key, username_input],
+        outputs=output_image
+    ).then(
+        fn=get_history,
+        inputs=username_input,
+        outputs=history_gallery
     )
     
     username_input.change(
         fn=get_history,
         inputs=username_input,
-        outputs=[history_gallery, history_gallery]
+        outputs=history_gallery
     )
     
     download_zip_btn.click(
         fn=download_history,
         outputs=download_zip_file
+    ).then(
+        lambda: gr.update(visible=True),
+        outputs=download_zip_file
     )
 
-# Add examples section
-gr.Examples(
-    examples=[
-        ["A cyberpunk cat wearing neon goggles, digital art style", "1:1"],
-        ["Mountain landscape at sunset with northern lights, photorealistic", "16:9"],
-        ["Steampunk airship floating above Victorian city, detailed illustration", "9:16"],
-        ["Minimalist logo of a bird made from geometric shapes, blue and white", "1:1"]
-    ],
-    inputs=[prompt_input, aspect_ratio],
-    label="💡 Example Prompts"
-)
-
 if __name__ == "__main__":
     demo.launch()