Update app.py

app.py

@@ -8,9 +8,27 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
+from huggingface_hub import HfFolder, Repository
+from pathlib import Path
 
-# Create permanent storage directory in the current working directory
-SAVE_DIR = os.path.join(os.getcwd(), "generated_images")
+# Set up Hugging Face repository for permanent storage
+REPO_ID = "YOUR_USERNAME/YOUR_SPACE_NAME"  # Replace with your actual repo ID
+REPO_TYPE = "space"
+HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space's secrets
+PERSISTENT_DIR = os.path.join(os.getcwd(), "persistent")
+
+# Initialize repository
+if HF_TOKEN:
+    HfFolder.save_token(HF_TOKEN)
+    repo = Repository(
+        local_dir=PERSISTENT_DIR,
+        clone_from=REPO_ID,
+        repo_type=REPO_TYPE,
+        use_auth_token=HF_TOKEN,
+    )
+
+# Create persistent storage directory
+SAVE_DIR = os.path.join(PERSISTENT_DIR, "generated_images")
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
@@ -25,6 +43,13 @@ pipeline = pipeline.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def git_push():
+    if HF_TOKEN:
+        repo.git_pull()
+        repo.git_add()
+        repo.git_commit("Update generated images")
+        repo.git_push()
+
 def save_generated_image(image, prompt):
     # Generate unique filename with timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -40,12 +65,19 @@ def save_generated_image(image, prompt):
     with open(metadata_file, "a", encoding="utf-8") as f:
         f.write(f"{filename}|{prompt}|{timestamp}\n")
     
+    # Push changes to Hugging Face
+    git_push()
+    
     return filepath
 
 def load_generated_images():
     if not os.path.exists(SAVE_DIR):
         return []
     
+    # Pull latest changes
+    if HF_TOKEN:
+        repo.git_pull()
+    
     # Load all images from the directory
     image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
                   if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]