Update app.py

app.py

@@ -9,12 +9,9 @@ import uuid
 from datetime import datetime
 import zipfile
 
-# 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")  # Fixed typo: BASE_dir -> BASE_DIR
 os.makedirs(HISTORY_DIR, exist_ok=True)
 ZIP_PATH = os.path.join(BASE_DIR, "image_history.zip")
 
@@ -26,8 +23,16 @@ if not os.path.exists(HISTORY_METADATA):
 
 def load_metadata():
     """Load the metadata file"""
-    with open(HISTORY_METADATA, "r") as f:
-        return json.load(f)
+    try:
+        with open(HISTORY_METADATA, "r") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, FileNotFoundError):
+        return []
+
+def save_metadata(metadata):
+    """Save metadata to file"""
+    with open(HISTORY_METADATA, "w") as f:
+        json.dump(metadata, f, indent=2)
 
 def save_image(image, username, prompt):
     """Save an image and record metadata"""
@@ -51,41 +56,98 @@ def get_history(username=None):
     metadata = load_metadata()
     
     # Filter by username if provided
-    if username:
-        metadata = [m for m in metadata if m["username"].lower() == username.lower()]
+    if username and username.strip():
+        metadata = [m for m in metadata if m["username"].lower() == username.lower().strip()]
     
-    # Prepare gallery data
-    history_data = []
+    # Prepare gallery data - return list of (image_path, label) tuples
+    gallery_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
+            label = f"{item['username']} - {item['timestamp'].split('T')[0]}"
+            gallery_data.append((item["path"], label))
+    
+    return gallery_data
 
-def download_history():
+def download_history(username=None):
     """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 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)
+    zip_path = os.path.join(BASE_DIR, f"{username or 'all'}_history.zip")
+    
+    with zipfile.ZipFile(zip_path, "w") as zipf:
+        metadata = load_metadata()
+        
+        # Filter by username if provided
+        if username and username.strip():
+            metadata = [m for m in metadata if m["username"].lower() == username.lower().strip()]
+        
+        # Add image files
+        for item in metadata:
+            if os.path.exists(item["path"]):
+                arcname = os.path.basename(item["path"])
+                zipf.write(item["path"], arcname=arcname)
         
         # Add metadata file
-        with open(HISTORY_METADATA, "rb") as meta_file:
-            zipf.write(meta_file, arcname="metadata.json")
+        metadata_str = json.dumps(metadata, indent=2)
+        zipf.writestr("metadata.json", metadata_str)
+    
+    return zip_path
+
+def generate_image(prompt, username):
+    """Generate image using Imagen 4.0 API"""
+    if not username or not username.strip():
+        raise gr.Error("Username cannot be empty")
+    
+    if not prompt or not prompt.strip():
+        raise gr.Error("Prompt cannot be empty")
     
-    return ZIP_PATH
+    # Get API key from environment variable
+    api_key = os.environ.get("BYTEZ_API_KEY")
+    if not api_key:
+        raise gr.Error("BYTEZ_API_KEY environment variable not set")
+    
+    headers = {
+        "Authorization": api_key,
+        "Content-Type": "application/json"
+    }
+    
+    payload = {
+        "text": prompt.strip(),
+        "aspect_ratio": "1:1"
+    }
+
+    try:
+        response = requests.post(
+            "https://api.bytez.com/models/v2/google/imagen-4.0-fast-generate-001",
+            headers=headers,
+            json=payload,
+            timeout=60
+        )
+        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("API returned no results")
+        
+        # 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.strip(), prompt.strip())
+        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:
+            try:
+                error_detail = e.response.json().get('error', {}).get('message', '')
+                if error_detail:
+                    error_msg += f"\n{error_detail}"
+            except:
+                error_msg += f"\nStatus {e.response.status_code}: {e.response.text[:200]}"
+        raise gr.Error(error_msg)
+    
+    except Exception as e:
+        raise gr.Error(f"Image generation failed: {str(e)}")
 
 with gr.Blocks(title="Imagen 4.0 Generator") as demo:
     gr.HTML("""
@@ -117,90 +179,31 @@ with gr.Blocks(title="Imagen 4.0 Generator") as demo:
             submit_btn = gr.Button("✨ Generate Image", variant="primary", size="lg")
             
             # 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, visible=True)
-            history_gallery = gr.Gallery(label="Your History", columns=3, preview=True, visible=True)
-            
-            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)
-    
-    # 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
+            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"
             )
             
-            # 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
+            download_zip_btn = gr.Button("📥 Download Your History ZIP", variant="secondary")
+            download_zip_file = gr.File(label="Download ZIP", visible=False)
         
-        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
+        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, object_fit="contain")
+
+    # Event handlers
     submit_btn.click(
         fn=generate_image,
-        inputs=[prompt_input, bytez_api_key, username_input],
+        inputs=[prompt_input, username_input],
         outputs=output_image
     ).then(
-        fn=update_history_display,
+        fn=get_history,
         inputs=username_input,
         outputs=history_gallery
     )
@@ -209,34 +212,16 @@ 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(
         fn=download_history,
+        inputs=username_input,
         outputs=download_zip_file
     ).then(
         lambda: gr.update(visible=True),
         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()