Upload 5 files

README.md

@@ -6,27 +6,40 @@ colorTo: purple
 sdk: gradio
 sdk_version: 4.19.2
 app_file: app.py
-pinned: true
+pinned: false
 license: mit
 ---
 
-# 🍌 NanoBanana Gemini Image Generator
+# 🍌 NanoBanana Gemini Image Generator - Version 3
 
-AI-powered image generation service using Google's Gemini 2.0 Flash model with Gradio UI and FastAPI REST endpoints.
+AI-powered image generation service using Google's Gemini 2.5 Flash Image Preview model with dataset repository integration and enhanced API features.
 
-## 🌟 Features
+## 🌟 Version 3 Features (Planned)
+
+### Dataset Integration
+- **Hugging Face Dataset Repository**: Automatic saving of generated images to dataset repository
+- **Metadata Tracking**: Store generation parameters and prompts with images
+- **Version History**: Track iterations and improvements
+
+### Enhanced API Features
+- **Batch Processing**: Generate multiple images in one request
+- **Webhooks**: Notify when generation completes
+- **Rate Limiting**: Control API usage
+- **API Keys**: User authentication and tracking
 
 ### Web Interface (Gradio)
-- **Generate**: Create images from text prompts using Gemini 2.0 Flash
+- **Generate**: Create images from text prompts using Gemini 2.5 Flash Image Preview
 - **Edit**: Modify existing images with text instructions
 - **Compose**: Combine multiple images into compositions
 - **History**: View recent generations with metadata
+- **Gallery**: Browse dataset repository images
 
 ### REST API (FastAPI)
 - Full REST API with automatic documentation
 - JSON request/response format
 - Base64 image encoding
 - Comprehensive error handling
+- Dataset repository integration
 
 ## 🚀 Quick Start
 
@@ -36,6 +49,10 @@ AI-powered image generation service using Google's Gemini 2.0 Flash model with G
    - In Hugging Face Spaces: Add `GEMINI_API_KEY` as a secret
    - Locally: Create `.env` file with `GEMINI_API_KEY=your_api_key_here`
 
+2. **Set Hugging Face Token (for dataset access)**
+   - Add `HF_TOKEN` as a secret in Spaces
+   - Or in `.env` file locally
+
 ### Access Points
 
 Once deployed:
@@ -67,10 +84,11 @@ GET /api/history?limit=10
 
 ## 🛠️ Technology Stack
 
-- **AI Model**: Google Gemini 2.0 Flash (Experimental)
+- **AI Model**: Google Gemini 2.5 Flash Image Preview
 - **Frontend**: Gradio 4.19.2
 - **Backend**: FastAPI
 - **Server**: Uvicorn (ASGI)
+- **Dataset**: Hugging Face Datasets
 - **Runtime**: Hugging Face Spaces (Gradio SDK)
 - **Python**: 3.10+
 
@@ -80,4 +98,4 @@ MIT License
 
 ---
 
-Made with ❤️ using Gradio, FastAPI, and Google Gemini
+Made with ❤️ using Gradio, FastAPI, Google Gemini, and Hugging Face

app.py

@@ -46,6 +46,24 @@ if GEMINI_API_KEY:
 else:
     logger.warning("GEMINI_API_KEY not found. Image generation will use placeholder images.")
 
+# Initialize Dataset Manager
+HF_TOKEN = os.getenv("HF_TOKEN")
+DATASET_REPO_ID = os.getenv("DATASET_REPO_ID")
+dataset_manager = None
+
+if HF_TOKEN and DATASET_REPO_ID:
+    try:
+        from dataset_manager import DatasetManager
+        dataset_manager = DatasetManager(DATASET_REPO_ID, HF_TOKEN)
+        logger.info(f"Dataset manager initialized for repository: {DATASET_REPO_ID}")
+    except Exception as e:
+        logger.warning(f"Could not initialize dataset manager: {e}")
+else:
+    if not HF_TOKEN:
+        logger.info("HF_TOKEN not set. Dataset saving feature disabled.")
+    if not DATASET_REPO_ID:
+        logger.info("DATASET_REPO_ID not set. Dataset saving feature disabled.")
+
 def generate_image_with_gemini(prompt: str, width: int = 1024, height: int = 1024, style: str = "Default") -> Image.Image:
     """
     Generate image using Gemini 2.5 Flash Image Preview (Nano Banana)
@@ -370,7 +388,9 @@ async def health_check():
 async def generate_image_api(
     prompt: str,
     size: str = "1024x1024",
-    style: str = "Default"
+    style: str = "Default",
+    save_to_dataset: bool = True,
+    dataset_folder: Optional[str] = None
 ):
     """Generate image via API using Gemini Nano Banana"""
     try:
@@ -380,18 +400,37 @@ async def generate_image_api(
         # Generate image
         image = generate_image_with_gemini(prompt, width, height, style)
 
-        # Save image
+        # Save image locally
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         filename = f"api_gen_{timestamp}.png"
         filepath = GENERATED_DIR / filename
         image.save(filepath)
 
+        # Save to dataset repository if enabled
+        dataset_url = None
+        if dataset_manager and save_to_dataset:
+            try:
+                metadata = {
+                    "style": style,
+                    "size": size,
+                    "model": MODEL_NAME if GEMINI_API_KEY else "placeholder",
+                    "generation_type": "text-to-image"
+                }
+                dataset_url = dataset_manager.save_image(
+                    image=image,
+                    prompt=prompt,
+                    folder_name=dataset_folder,
+                    metadata=metadata
+                )
+            except Exception as dataset_error:
+                logger.error(f"Failed to save to dataset: {dataset_error}")
+
         # Convert to base64
         buffer = BytesIO()
         image.save(buffer, format="PNG")
         img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
 
-        return JSONResponse(content={
+        response_data = {
             "success": True,
             "filename": filename,
             "prompt": prompt,
@@ -399,7 +438,12 @@ async def generate_image_api(
             "style": style,
             "model": MODEL_NAME if GEMINI_API_KEY else "placeholder",
             "image_base64": img_base64
-        })
+        }
+
+        if dataset_url:
+            response_data["dataset_url"] = dataset_url
+
+        return JSONResponse(content=response_data)
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -460,8 +504,39 @@ async def get_generation_history(limit: int = 10):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.get("/api/dataset/folders")
+async def get_dataset_folders():
+    """Get list of folders in dataset repository"""
+    if not dataset_manager:
+        return JSONResponse(
+            status_code=503,
+            content={"error": "Dataset manager not configured"}
+        )
+
+    try:
+        folders = dataset_manager.get_folders()
+        return JSONResponse(content={"folders": folders, "count": len(folders)})
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/dataset/images/{folder_name}")
+async def get_dataset_images(folder_name: str):
+    """Get list of images in a specific dataset folder"""
+    if not dataset_manager:
+        return JSONResponse(
+            status_code=503,
+            content={"error": "Dataset manager not configured"}
+        )
+
+    try:
+        images = dataset_manager.get_images_in_folder(folder_name)
+        return JSONResponse(content={"images": images, "count": len(images), "folder": folder_name})
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 # Gradio Interface functions
-def gradio_generate(prompt: str, size: str, style: str, quality: str, negative_prompt: str):
+def gradio_generate(prompt: str, size: str, style: str, quality: str, negative_prompt: str,
+                   save_to_dataset: bool = True, dataset_folder: str = ""):
     """Generate image through Gradio interface using Nano Banana"""
     try:
         if not prompt:
@@ -485,7 +560,7 @@ def gradio_generate(prompt: str, size: str, style: str, quality: str, negative_p
         # Generate image using Gemini
         image = generate_image_with_gemini(enhanced_prompt, width, height, style)
 
-        # Save image
+        # Save image locally
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         filename = f"gradio_gen_{timestamp}.png"
         filepath = GENERATED_DIR / filename
@@ -497,6 +572,33 @@ def gradio_generate(prompt: str, size: str, style: str, quality: str, negative_p
         else:
             status += f"\n🎨 Model: {MODEL_NAME}"
 
+        # Save to dataset if enabled
+        if dataset_manager and save_to_dataset:
+            try:
+                metadata = {
+                    "style": style,
+                    "quality": quality,
+                    "size": f"{width}x{height}",
+                    "negative_prompt": negative_prompt,
+                    "model": MODEL_NAME if GEMINI_API_KEY else "placeholder",
+                    "generation_type": "text-to-image"
+                }
+
+                # Use provided folder or None (will default to date)
+                folder_name = dataset_folder if dataset_folder.strip() else None
+
+                dataset_url = dataset_manager.save_image(
+                    image=image,
+                    prompt=prompt,
+                    folder_name=folder_name,
+                    metadata=metadata
+                )
+
+                if dataset_url:
+                    status += f"\n📁 Saved to dataset: {folder_name or datetime.now().strftime('%Y_%m_%d')}"
+            except Exception as dataset_error:
+                status += f"\n⚠️ Dataset save failed: {str(dataset_error)}"
+
         return image, status
 
     except Exception as e:
@@ -573,6 +675,40 @@ def gradio_compose(images, compose_prompt):
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
 
+def get_dataset_folders_list():
+    """Get list of folders for dropdown"""
+    if not dataset_manager:
+        return []
+    try:
+        return dataset_manager.get_folders()
+    except:
+        return []
+
+def load_dataset_gallery(folder_name: str):
+    """Load images from dataset folder for gallery"""
+    if not dataset_manager or not folder_name:
+        return []
+
+    try:
+        images = dataset_manager.get_images_in_folder(folder_name)
+        gallery_items = []
+
+        # Load first few images for preview
+        for img_info in images[:20]:  # Limit to 20 for performance
+            try:
+                # Download and display image from URL
+                import requests
+                response = requests.get(img_info["url"], timeout=10)
+                if response.status_code == 200:
+                    img = Image.open(BytesIO(response.content))
+                    gallery_items.append(img)
+            except:
+                continue
+
+        return gallery_items
+    except:
+        return []
+
 # Create Gradio interface
 with gr.Blocks(title="NanoBanana Gemini Image Generator", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
@@ -649,6 +785,23 @@ with gr.Blocks(title="NanoBanana Gemini Image Generator", theme=gr.themes.Soft()
                         lines=2
                     )
 
+                    # Dataset save options
+                    with gr.Accordion("📁 Dataset Options", open=False):
+                        gen_save_dataset = gr.Checkbox(
+                            label="Save to Dataset Repository",
+                            value=True if dataset_manager else False,
+                            interactive=bool(dataset_manager)
+                        )
+                        gen_dataset_folder = gr.Textbox(
+                            label="Folder Name (leave empty for date-based)",
+                            placeholder="e.g., 'portraits' or leave empty for YYYY_MM_DD",
+                            value="",
+                            interactive=bool(dataset_manager)
+                        )
+
+                        if not dataset_manager:
+                            gr.Markdown("⚠️ Dataset saving disabled. Set HF_TOKEN and DATASET_REPO_ID in environment.")
+
                     gen_button = gr.Button("🚀 Generate with Nano Banana", variant="primary", size="lg")
 
                 with gr.Column():
@@ -671,7 +824,7 @@ with gr.Blocks(title="NanoBanana Gemini Image Generator", theme=gr.themes.Soft()
 
             gen_button.click(
                 fn=gradio_generate,
-                inputs=[gen_prompt, gen_size, gen_style, gen_quality, gen_negative],
+                inputs=[gen_prompt, gen_size, gen_style, gen_quality, gen_negative, gen_save_dataset, gen_dataset_folder],
                 outputs=[gen_output, gen_status]
             )
 
@@ -771,6 +924,62 @@ with gr.Blocks(title="NanoBanana Gemini Image Generator", theme=gr.themes.Soft()
             # Auto-load history on tab open
             demo.load(fn=get_history, outputs=history_display)
 
+        # Dataset Gallery Tab
+        if dataset_manager:
+            with gr.Tab("🖼️ Dataset Gallery"):
+                gr.Markdown(
+                    f"""
+                    ### Browse Generated Images
+                    View images saved to the dataset repository: **{DATASET_REPO_ID}**
+
+                    [View on Hugging Face](https://huggingface.co/datasets/{DATASET_REPO_ID})
+                    """
+                )
+
+                with gr.Row():
+                    folder_dropdown = gr.Dropdown(
+                        label="Select Folder",
+                        choices=get_dataset_folders_list(),
+                        value=None,
+                        interactive=True
+                    )
+                    refresh_folders_btn = gr.Button("🔄 Refresh Folders")
+
+                gallery = gr.Gallery(
+                    label="Images",
+                    show_label=True,
+                    elem_id="gallery",
+                    columns=4,
+                    rows=3,
+                    object_fit="contain",
+                    height="auto"
+                )
+
+                def refresh_folders():
+                    folders = get_dataset_folders_list()
+                    return gr.Dropdown(choices=folders, value=folders[0] if folders else None)
+
+                refresh_folders_btn.click(
+                    fn=refresh_folders,
+                    outputs=folder_dropdown
+                )
+
+                folder_dropdown.change(
+                    fn=load_dataset_gallery,
+                    inputs=folder_dropdown,
+                    outputs=gallery
+                )
+
+                # Load first folder on tab open
+                demo.load(
+                    fn=lambda: get_dataset_folders_list()[0] if get_dataset_folders_list() else None,
+                    outputs=folder_dropdown
+                ).then(
+                    fn=load_dataset_gallery,
+                    inputs=folder_dropdown,
+                    outputs=gallery
+                )
+
     # Footer
     gr.Markdown(
         """

dataset_manager.py

@@ -0,0 +1,294 @@
+"""
+Dataset Manager for NanoBanana Image Generator
+Handles saving generated images to Hugging Face Dataset Repository
+"""
+
+import os
+import json
+import logging
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+from pathlib import Path
+from PIL import Image
+from io import BytesIO
+import tempfile
+
+from huggingface_hub import HfApi, upload_file, upload_folder, create_repo, list_repo_files
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetManager:
+    """
+    Manages image uploads to Hugging Face Dataset Repository
+    with custom folder structure support
+    """
+
+    def __init__(self, repo_id: str, token: str):
+        """
+        Initialize Dataset Manager
+
+        Args:
+            repo_id: Hugging Face dataset repository ID (e.g., "username/dataset-name")
+            token: Hugging Face API token with write access
+        """
+        self.repo_id = repo_id
+        self.token = token
+        self.api = HfApi(token=token)
+
+        # Ensure repository exists
+        self._ensure_repo_exists()
+
+        logger.info(f"DatasetManager initialized for repository: {repo_id}")
+
+    def _ensure_repo_exists(self):
+        """Create dataset repository if it doesn't exist"""
+        try:
+            # Try to get repository info
+            self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")
+            logger.info(f"Dataset repository {self.repo_id} already exists")
+        except Exception as e:
+            # Repository doesn't exist, create it
+            try:
+                self.api.create_repo(
+                    repo_id=self.repo_id,
+                    repo_type="dataset",
+                    private=False,
+                    exist_ok=True
+                )
+                logger.info(f"Created new dataset repository: {self.repo_id}")
+
+                # Create initial README
+                readme_content = f"""# NanoBanana Generated Images Dataset
+
+This dataset contains images generated by the NanoBanana Gemini Image Generator.
+
+## Structure
+- `images/` - Generated images organized by date or custom folders
+  - `YYYY_MM_DD/` - Images generated on specific dates
+  - Custom folders as specified during generation
+
+## Metadata
+Each image is accompanied by a JSON metadata file containing:
+- Generation prompt
+- Model used
+- Generation parameters
+- Timestamp
+
+## Usage
+You can load this dataset using the Hugging Face `datasets` library:
+
+```python
+from datasets import load_dataset
+dataset = load_dataset("{self.repo_id}")
+```
+
+---
+Generated with 🍌 NanoBanana Image Generator
+"""
+                self._upload_readme(readme_content)
+            except Exception as create_error:
+                logger.warning(f"Could not create repository: {create_error}")
+
+    def _upload_readme(self, content: str):
+        """Upload README.md to repository"""
+        try:
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
+                f.write(content)
+                temp_path = f.name
+
+            self.api.upload_file(
+                path_or_fileobj=temp_path,
+                path_in_repo="README.md",
+                repo_id=self.repo_id,
+                repo_type="dataset"
+            )
+
+            # Clean up temp file
+            os.unlink(temp_path)
+        except Exception as e:
+            logger.error(f"Failed to upload README: {e}")
+
+    def save_image(
+        self,
+        image: Image.Image,
+        prompt: str,
+        folder_name: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> Optional[str]:
+        """
+        Save image to dataset repository
+
+        Args:
+            image: PIL Image object to save
+            prompt: Generation prompt used
+            folder_name: Optional folder name. If None, uses YYYY_MM_DD format
+            metadata: Optional additional metadata to save
+
+        Returns:
+            URL to the uploaded image in the dataset repository, or None if failed
+        """
+        try:
+            # Determine folder name
+            if not folder_name:
+                folder_name = datetime.now().strftime("%Y_%m_%d")
+
+            # Generate unique filename
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]  # Milliseconds
+            image_filename = f"image_{timestamp}.png"
+            metadata_filename = f"image_{timestamp}.json"
+
+            # Paths in repository
+            image_path_in_repo = f"images/{folder_name}/{image_filename}"
+            metadata_path_in_repo = f"images/{folder_name}/{metadata_filename}"
+
+            # Create metadata
+            image_metadata = {
+                "prompt": prompt,
+                "timestamp": datetime.now().isoformat(),
+                "folder": folder_name,
+                "filename": image_filename,
+                "size": {
+                    "width": image.width,
+                    "height": image.height
+                }
+            }
+
+            # Add additional metadata if provided
+            if metadata:
+                image_metadata.update(metadata)
+
+            # Save image to temporary file
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as img_temp:
+                image.save(img_temp, format='PNG')
+                img_temp_path = img_temp.name
+
+            # Save metadata to temporary file
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as meta_temp:
+                json.dump(image_metadata, meta_temp, indent=2)
+                meta_temp_path = meta_temp.name
+
+            # Upload image
+            self.api.upload_file(
+                path_or_fileobj=img_temp_path,
+                path_in_repo=image_path_in_repo,
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                commit_message=f"Add generated image: {image_filename}"
+            )
+
+            # Upload metadata
+            self.api.upload_file(
+                path_or_fileobj=meta_temp_path,
+                path_in_repo=metadata_path_in_repo,
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                commit_message=f"Add metadata for: {image_filename}"
+            )
+
+            # Clean up temp files
+            os.unlink(img_temp_path)
+            os.unlink(meta_temp_path)
+
+            # Return URL to the image
+            dataset_url = f"https://huggingface.co/datasets/{self.repo_id}/blob/main/{image_path_in_repo}"
+
+            logger.info(f"Successfully saved image to dataset: {image_path_in_repo}")
+            return dataset_url
+
+        except Exception as e:
+            logger.error(f"Failed to save image to dataset: {e}")
+            return None
+
+    def get_folders(self) -> List[str]:
+        """
+        Get list of folders in the dataset repository
+
+        Returns:
+            List of folder names
+        """
+        try:
+            files = self.api.list_repo_files(
+                repo_id=self.repo_id,
+                repo_type="dataset"
+            )
+
+            # Extract folder names from paths
+            folders = set()
+            for file in files:
+                if file.startswith("images/"):
+                    parts = file.split("/")
+                    if len(parts) > 1:
+                        folder = parts[1]
+                        folders.add(folder)
+
+            return sorted(list(folders))
+
+        except Exception as e:
+            logger.error(f"Failed to get folders: {e}")
+            return []
+
+    def get_images_in_folder(self, folder_name: str) -> List[Dict[str, str]]:
+        """
+        Get list of images in a specific folder
+
+        Args:
+            folder_name: Name of the folder
+
+        Returns:
+            List of dictionaries containing image info
+        """
+        try:
+            files = self.api.list_repo_files(
+                repo_id=self.repo_id,
+                repo_type="dataset"
+            )
+
+            images = []
+            folder_prefix = f"images/{folder_name}/"
+
+            for file in files:
+                if file.startswith(folder_prefix) and file.endswith(".png"):
+                    # Get corresponding metadata file
+                    metadata_file = file.replace(".png", ".json")
+
+                    image_info = {
+                        "filename": os.path.basename(file),
+                        "path": file,
+                        "url": f"https://huggingface.co/datasets/{self.repo_id}/resolve/main/{file}",
+                        "metadata_url": f"https://huggingface.co/datasets/{self.repo_id}/resolve/main/{metadata_file}"
+                    }
+                    images.append(image_info)
+
+            return sorted(images, key=lambda x: x["filename"], reverse=True)
+
+        except Exception as e:
+            logger.error(f"Failed to get images in folder {folder_name}: {e}")
+            return []
+
+    def batch_save_images(
+        self,
+        images_data: List[Dict[str, Any]],
+        folder_name: Optional[str] = None
+    ) -> List[Optional[str]]:
+        """
+        Save multiple images in batch
+
+        Args:
+            images_data: List of dicts with 'image', 'prompt', and optional 'metadata'
+            folder_name: Optional folder name for all images
+
+        Returns:
+            List of URLs for uploaded images
+        """
+        urls = []
+        for data in images_data:
+            url = self.save_image(
+                image=data['image'],
+                prompt=data['prompt'],
+                folder_name=folder_name,
+                metadata=data.get('metadata')
+            )
+            urls.append(url)
+
+        return urls

requirements.txt

@@ -13,6 +13,8 @@ numpy>=1.24.0
 # Utilities
 python-dotenv>=1.0.0
 aiofiles>=23.2.1
+requests>=2.28.0
 
-# For image generation via nanobanana MCP
-huggingface_hub>=0.20.0
+# Hugging Face integration
+huggingface_hub>=0.20.0
+datasets>=2.14.0