Update app.py

app.py

@@ -2,113 +2,90 @@ import os
 import base64
 import json
 import requests
-import gradio as gr
+import gr
 from io import BytesIO
 from PIL import Image
 import uuid
 from datetime import datetime
 import zipfile
 
-bytez_api_key = os.getenv("BYTEZ_API_KEY")
+# Load environment variables
+bytez_api_key = os.getenv("BYTEZ_API_KEY") or input("Enter Bytez API Key (press Enter for default): ")
 
 # Persistent storage setup
 BASE_DIR = os.path.abspath(os.getcwd())
-HISTORY_DIR = os.path.join(BASE_DIR, "image_history")
+HISTORY_DIR = os.path.join(BASE_dir, "image_history")
 os.makedirs(HISTORY_DIR, exist_ok=True)
 ZIP_PATH = os.path.join(BASE_DIR, "image_history.zip")
 
-# Store metadata
+# Metadata file
 HISTORY_METADATA = os.path.join(HISTORY_DIR, "metadata.json")
 if not os.path.exists(HISTORY_METADATA):
     with open(HISTORY_METADATA, "w") as f:
         json.dump([], f)
 
 def load_metadata():
+    """Load the metadata file"""
     with open(HISTORY_METADATA, "r") as f:
         return json.load(f)
 
-def save_metadata(metadata):
-    with open(HISTORY_METADATA, "w") as f:
-        json.dump(metadata, f, indent=2)
-
-def save_image(image: Image.Image, username, prompt):
+def save_image(image, username, prompt):
+    """Save an image and record metadata"""
     filename = os.path.join(HISTORY_DIR, f"{uuid.uuid4()}.png")
     image.save(filename)
+    
     metadata = load_metadata()
+    timestamp = datetime.now().isoformat()
     metadata.append({
         "path": filename,
         "username": username,
         "prompt": prompt,
-        "timestamp": datetime.now().isoformat()
+        "timestamp": timestamp
     })
+    
     save_metadata(metadata)
     return filename
 
 def get_history(username=None):
+    """Get image history by username"""
     metadata = load_metadata()
+    
+    # Filter by username if provided
     if username:
-        metadata = [m for m in metadata if m["username"] == username]
-    else:
-        return [], []
-    return [(m["path"], f"{m['username']} - {m['timestamp']}") for m in metadata]
+        metadata = [m for m in metadata if m["username"].lower() == username.lower()]
+    
+    # Prepare gallery data
+    history_data = []
+    for item in metadata:
+        if os.path.exists(item["path"]):
+            with open(item["path"], "rb") as img_file:
+                img = Image.open(img_file)
+                img.load()  # Required for getting size
+                history_data.append({
+                    "image": img,
+                    "username": item["username"],
+                    "timestamp": item["timestamp"]
+                })
+    
+    return history_data
 
 def download_history():
+    """Download history as ZIP file"""
     with zipfile.ZipFile(ZIP_PATH, "w") as zipf:
+        # Add all image files
         for root, dirs, files in os.walk(HISTORY_DIR):
             for file in files:
-                if file != "metadata.json":
-                    zipf.write(os.path.join(root, file), arcname=file)
-        zipf.write(HISTORY_METADATA, arcname="metadata.json")
-    return ZIP_PATH
-
-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 gr.Error("No API key provided. Set BYTEZ_API_KEY environment variable or enter in UI.")
-    
-    headers = {
-        "Authorization": api_key,
-        "Content-Type": "application/json"
-    }
-    
-    payload = {
-        "text": prompt,
-        #"aspect_ratio": "1:1"  # Fixed to 1:1 ratio as requested
-    }
-
-    try:
-        response = requests.post(
-            "https://api.bytez.com/models/v2/google/imagen-4.0-fast-generate-001",
-            headers=headers,
-            json=payload,
-            timeout=45
-        )
-        response.raise_for_status()
-        
-        result = response.json()
-        if not result.get("results") or not isinstance(result["results"], list) or len(result["results"]) == 0:
-            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))
-        
-        # Save to history
-        saved_path = save_image(img, username, prompt)
+                if not file.endswith(".json"):
+                    file_path = os.path.join(root, file)
+                    if os.path.exists(file_path):
+                        with open(file_path, "rb") as img_file:
+                            zipf.write(img_file, arcname=file)
         
-        return saved_path
+        # Add metadata file
+        with open(HISTORY_METADATA, "rb") as meta_file:
+            zipf.write(meta_file, arcname="metadata.json")
     
-    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 gr.Error(error_msg)
-    except Exception as e:
-        raise gr.Error(f"Image generation failed: {str(e)}")
+    return ZIP_PATH
 
 with gr.Blocks(title="Imagen 4.0 Generator") as demo:
     gr.HTML("""
@@ -122,49 +99,108 @@ with gr.Blocks(title="Imagen 4.0 Generator") as demo:
         </div>
     </div>
     """)
-
+    
     with gr.Row():
         with gr.Column(scale=1):
             username_input = gr.Textbox(
-                lines=1, 
-                placeholder="Enter your username", 
+                lines=1,
+                placeholder="Enter your username",
                 label="Username",
                 info="Required for saving history"
             )
             prompt_input = gr.Textbox(
-                lines=4, 
-                placeholder="A Raspberry Pi SBC on a wooden desk with electronic components", 
+                lines=4,
+                placeholder="A Raspberry Pi SBC on a wooden desk with electronic components",
                 label="Prompt",
                 info="Be descriptive for best results"
             )
             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"
-            )
-
+            # Example prompts
+            with gr.Accordion("Example Prompts", open=False):
+                example_prompts = [
+                    "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"
+                ]
+                gr.Button("🔄 Refresh Examples").click(
+                    fn=lambda: example_prompts,
+                    outputs=gr.HighlightedText(label="Example Prompts")
+                )
+            
         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)
+            output_image = gr.Image(label="Generated Image", height=512, visible=True)
+            history_gallery = gr.Gallery(label="Your History", columns=3, preview=True, visible=True)
             
-            with gr.Row():
-                download_zip_btn = gr.Button("📥 Download History ZIP")
+            with gr.Row(visible=False):
+                download_zip_btn = gr.Button("📥 Download History ZIP", variant="secondary")
                 download_zip_file = gr.File(label="Download ZIP", visible=False)
-
-    # Event handlers
+    
+    # Generate image function
+    def generate_image(prompt, bytez_api_key, username):
+        if not username.strip():
+            raise gr.Error("Username cannot be empty")
+        
+        # Get API key from Bytez environment variable as fallback
+        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 gr.Error("No API key provided. Set BYTEZ_API_KEY environment variable or enter in UI.")
+        
+        # Create headers with authentication
+        headers = {
+            "Authorization": api_key,
+            "Content-Type": "application/json"
+        }
+        
+        # Prepare payload
+        payload = {
+            "text": prompt
+        }
+        
+        # Send request to Bytez API
+        try:
+            response = requests.post(
+                "https://api.bytez.com/models/v2/google/imagen-4.0-fast-generate-001",
+                headers=headers,
+                json=payload,
+                timeout=45  # Increased timeout for image generation
+            )
+            
+            # Check response
+            response.raise_for_status()
+            
+            # Parse JSON response
+            result = response.json()
+            
+            # Extract image data
+            if not result.get("results") or not isinstance(result["results"], list) or len(result["results"]) == 0:
+                raise gr.Error("No results returned from API")
+            
+            img_data = base64.b64decode(result["results"][0]["base64"])
+            img = Image.open(BytesIO(img_data))
+            
+            # Save image and return path
+            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 gr.Error(error_msg)
+        
+        except Exception as e:
+            raise gr.Error(f"Image generation failed: {str(e)}")
+    
+    # Connect functions
     submit_btn.click(
         fn=generate_image,
         inputs=[prompt_input, bytez_api_key, username_input],
         outputs=output_image
     ).then(
-        fn=get_history,
+        fn=update_history_display,
         inputs=username_input,
         outputs=history_gallery
     )
@@ -173,6 +209,9 @@ with gr.Blocks(title="Imagen 4.0 Generator") as demo:
         fn=get_history,
         inputs=username_input,
         outputs=history_gallery
+    ).then(
+        fn=update_image_visibility,
+        outputs=[output_image, history_gallery, download_zip_btn, download_zip_file]
     )
     
     download_zip_btn.click(
@@ -183,5 +222,21 @@ with gr.Blocks(title="Imagen 4.0 Generator") as demo:
         outputs=download_zip_file
     )
 
+def update_history_display(username):
+    """Update the history gallery display"""
+    try:
+        history_items = get_history(username)
+        return [item for item in history_items if item and os.path.exists(item["path"])]
+    except Exception as e:
+        raise gr.Error(f"Failed to load history: {str(e)}")
+
+def update_image_visibility(output_image, history_gallery, download_zip_btn, download_zip_file):
+    """Update visibility of components when needed"""
+    # Show image and hide gallery initially
+    if output_image is None:
+        return None, [], None, None
+    else:
+        return output_image, [], download_zip_btn, download_zip_file
+
 if __name__ == "__main__":
     demo.launch()