:clown_face: silly errors - dataset persistent storage

gradio_neutral_input_func.py

@@ -6,12 +6,18 @@ import json
 import uuid
 import os
 from stable_diffusion_demo import StableDiffusion
+from datasets import Dataset, Features, Value, Image as HFImage, load_dataset
+import tempfile
 
 # Setup directories
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
 IMAGE_DIR = os.path.join(BASE_DIR, "neutral_images_storage")
 os.makedirs(IMAGE_DIR, exist_ok=True)
 
+# HuggingFace dataset configuration
+DATASET_REPO = "willsh1997/neutral-sd-outputs"
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+
 def generate_image():
     """Generate a neutral image using Stable Diffusion"""
     generated_image = StableDiffusion(
@@ -25,56 +31,193 @@ def generate_image():
     )
     return generated_image
 
+def load_dataset_from_hf():
+    """Load dataset from HuggingFace Hub"""
+    try:
+        dataset = load_dataset(DATASET_REPO, split="train")
+        return dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        # Return empty dataset with correct schema if repo doesn't exist
+        return Dataset.from_dict({
+            "image": [],
+            "description": [],
+            "uuid": []
+        }).cast_column("image", HFImage())
+
+def save_to_hf_dataset(image, description):
+    """Save new image and description to HuggingFace dataset"""
+    try:
+        # Generate UUID for the new entry
+        image_id = str(uuid.uuid4())
+        
+        # Load existing dataset
+        try:
+            existing_dataset = load_dataset(DATASET_REPO, split="train")
+        except:
+            # Create empty dataset if it doesn't exist
+            existing_dataset = Dataset.from_dict({
+                "image": [],
+                "description": [],
+                "uuid": []
+            }).cast_column("image", HFImage())
+        
+        # Create temporary file for the image
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
+            image.save(tmp_file.name, format='PNG')
+            
+            # Create new entry
+            new_entry = {
+                "image": [tmp_file.name],
+                "description": [description],
+                "uuid": [image_id]
+            }
+            
+            # Create new dataset from the entry
+            new_dataset = Dataset.from_dict(new_entry).cast_column("image", HFImage())
+            
+            # Concatenate with existing dataset
+            if len(existing_dataset) > 0:
+                combined_dataset = existing_dataset.concatenate(new_dataset)
+            else:
+                combined_dataset = new_dataset
+            
+            # Push to HuggingFace Hub
+            combined_dataset.push_to_hub(DATASET_REPO, private=False, token=HF_TOKEN)
+            
+            # Clean up temporary file
+            os.unlink(tmp_file.name)
+            
+        return True, "Successfully saved to HuggingFace dataset!"
+        
+    except Exception as e:
+        return False, f"Error saving to HuggingFace: {str(e)}"
+
 def save_image_and_description(image, description):
-    """Save the generated image and its description"""
+    """Save the generated image and its description to HuggingFace dataset"""
     if image is None:
         return "No image to save!", None, None
     
     if not description:
         return "Please provide a description!", None, None
+    
+    # Save to HuggingFace dataset
+    success, message = save_to_hf_dataset(image, description)
+    
+    if success:
+        # Also save locally for backup/caching
+        try:
+            image_id = uuid.uuid4()
+            save_path = os.path.join(IMAGE_DIR, f"{image_id}.png")
+            json_path = os.path.join(IMAGE_DIR, f"{image_id}.json")
+            
+            image.save(save_path)
+            desc_json = {"description": description}
+            with open(json_path, "w") as f:
+                json.dump(desc_json, f)
+        except:
+            pass  # Local save is just backup, don't fail if it doesn't work
         
+        return message, None, load_previous_examples()
+    else:
+        return message, None, None
+
+def load_previous_examples():
+    """Load examples from HuggingFace dataset"""
     try:
-        image_id = uuid.uuid4()
-        save_path = os.path.join(IMAGE_DIR, f"{image_id}.png")
-        json_path = os.path.join(IMAGE_DIR, f"{image_id}.json")
+        dataset = load_dataset_from_hf()
+        examples = []
         
-        # Save image
-        image.save(save_path)
+        # Convert dataset to gallery format
+        for item in dataset:
+            if item['image'] is not None and item['description']:
+                examples.append((item['image'], item['description']))
         
-        # Save description
-        desc_json = {"description": description}
-        with open(json_path, "w") as f:
-            json.dump(desc_json, f)
+        return examples
         
-        # Return success message, clear the image output, and return updated gallery
-        return "Saved successfully!", None, load_previous_examples()
     except Exception as e:
-        return f"Error saving: {str(e)}", None, None
+        print(f"Error loading examples from HuggingFace: {e}")
+        # Fallback to local examples
+        return load_local_examples()
 
-def load_previous_examples():
-    """Load all previously saved images and descriptions"""
+def load_local_examples():
+    """Fallback: Load examples from local storage"""
     examples = []
-    for file in os.listdir(IMAGE_DIR):
-        if file.endswith(".png"):
-            image_id = file.replace(".png", "")
-            image_path = os.path.join(IMAGE_DIR, f"{image_id}.png")
-            json_path = os.path.join(IMAGE_DIR, f"{image_id}.json")
-            
-            if os.path.exists(json_path):
-                image = Image.open(image_path)
-                with open(json_path, "r") as f:
-                    desc = json.load(f)["description"]
-                examples.append((image, desc))
+    try:
+        for file in os.listdir(IMAGE_DIR):
+            if file.endswith(".png"):
+                image_id = file.replace(".png", "")
+                image_path = os.path.join(IMAGE_DIR, f"{image_id}.png")
+                json_path = os.path.join(IMAGE_DIR, f"{image_id}.json")
+                
+                if os.path.exists(json_path):
+                    image = Image.open(image_path)
+                    with open(json_path, "r") as f:
+                        desc = json.load(f)["description"]
+                    examples.append((image, desc))
+    except Exception as e:
+        print(f"Error loading local examples: {e}")
+    
     return examples
 
+def create_initial_dataset():
+    """Create initial dataset from local files if HF dataset doesn't exist"""
+    try:
+        # Check if we have local files to upload
+        local_examples = load_local_examples()
+        if not local_examples:
+            return
+        
+        # Try to load existing dataset
+        try:
+            existing_dataset = load_dataset(DATASET_REPO, split="train")
+            if len(existing_dataset) > 0:
+                return  # Dataset already exists with data
+        except:
+            pass  # Dataset doesn't exist, we'll create it
+        
+        # Create dataset from local files
+        images = []
+        descriptions = []
+        uuids = []
+        
+        for file in os.listdir(IMAGE_DIR):
+            if file.endswith(".png"):
+                image_id = file.replace(".png", "")
+                image_path = os.path.join(IMAGE_DIR, f"{image_id}.png")
+                json_path = os.path.join(IMAGE_DIR, f"{image_id}.json")
+                
+                if os.path.exists(json_path):
+                    with open(json_path, "r") as f:
+                        desc = json.load(f)["description"]
+                    
+                    images.append(image_path)
+                    descriptions.append(desc)
+                    uuids.append(image_id)
+        
+        if images:
+            # Create dataset
+            dataset_dict = {
+                "image": images,
+                "description": descriptions,
+                "uuid": uuids
+            }
+            
+            dataset = Dataset.from_dict(dataset_dict).cast_column("image", HFImage())
+            dataset.push_to_hub(DATASET_REPO, private=False)
+            print(f"Uploaded {len(images)} images to HuggingFace dataset")
+            
+    except Exception as e:
+        print(f"Error creating initial dataset: {e}")
+
 # Create the Gradio interface
 with gr.Blocks(title="Neutral Image App") as demo:
     gr.Markdown("# Neutral Image App")
+    gr.Markdown(f"*Images are saved to HuggingFace dataset: [{DATASET_REPO}](https://huggingface.co/datasets/{DATASET_REPO})*")
     
     with gr.Row():
         with gr.Column():
             generate_btn = gr.Button("Generate Image")
-            # Disable image upload by setting interactive=False
             image_output = gr.Image(type="pil", label="Generated Image", interactive=False)
             description_input = gr.Textbox(label="Describe the image", lines=3)
             save_btn = gr.Button("Save Image and Description")
@@ -82,10 +225,11 @@ with gr.Blocks(title="Neutral Image App") as demo:
     
     with gr.Accordion("Previous Examples", open=False):
         gallery = gr.Gallery(
-            label="Previous Images",
+            label="Previous Images from HuggingFace Dataset",
             show_label=True,
             elem_id="gallery"
-        )#.style(grid=2, height="auto")
+        )
+        refresh_btn = gr.Button("Refresh Gallery")
     
     # Set up event handlers
     generate_btn.click(
@@ -93,11 +237,15 @@ with gr.Blocks(title="Neutral Image App") as demo:
         outputs=[image_output]
     )
     
-    # Updated to include gallery refresh in outputs
     save_btn.click(
         fn=save_image_and_description,
         inputs=[image_output, description_input],
-        outputs=[status_output, image_output, gallery]  # Added gallery to outputs
+        outputs=[status_output, image_output, gallery]
+    )
+    
+    refresh_btn.click(
+        fn=load_previous_examples,
+        outputs=[gallery]
     )
     
     # Load previous examples on startup
@@ -108,4 +256,6 @@ with gr.Blocks(title="Neutral Image App") as demo:
 
 # Launch the app
 if __name__ == "__main__":
+    # Create initial dataset from local files if needed
+    create_initial_dataset()
     demo.launch()