Update app.py

app.py

@@ -11,19 +11,18 @@ import gradio as gr
 import torch
 from diffusers import FluxPipeline
 from PIL import Image
-# Setup and initialization code
-cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-# 이 부분을 수정
+
 
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-# Change gallery path to user's home directory for persistence
-# Setup and initialization code
-# Hugging Face Spaces permanent storage path
-gallery_path = "/data/gallery"  # 영구 저장을 위한 Spaces /data 디렉토리 사용
+# Use PERSISTENT_DIR environment variable for Spaces
+PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
+gallery_path = path.join(PERSISTENT_DIR, "gallery")
+
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
+
 torch.backends.cuda.matmul.allow_tf32 = True
 
 # Create gallery directory if it doesn't exist
@@ -151,36 +150,44 @@ footer {display: none !important}
     max-width: 100% !important;
 }
 """
-
 def save_image(image):
     """Save the generated image and return the path"""
-    # Ensure gallery directory exists
-    os.makedirs(gallery_path, exist_ok=True)
-    
-    # Generate unique filename with timestamp and random suffix
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    random_suffix = os.urandom(4).hex()
-    filename = f"generated_{timestamp}_{random_suffix}.png"
-    filepath = os.path.join(gallery_path, filename)
-    
     try:
-        if isinstance(image, Image.Image):
-            # Save with maximum quality
-            image.save(filepath, "PNG", quality=100)
-        else:
-            image = Image.fromarray(image)
-            image.save(filepath, "PNG", quality=100)
+        # Ensure gallery directory exists
+        if not os.path.exists(gallery_path):
+            try:
+                os.makedirs(gallery_path, exist_ok=True)
+            except Exception as e:
+                print(f"Failed to create gallery directory: {str(e)}")
+                return None
         
-        # Verify the file was saved correctly
-        if not os.path.exists(filepath):
-            print(f"Warning: Failed to verify saved image at {filepath}")
+        # Generate unique filename with timestamp and random suffix
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        random_suffix = os.urandom(4).hex()
+        filename = f"generated_{timestamp}_{random_suffix}.png"
+        filepath = os.path.join(gallery_path, filename)
+        
+        try:
+            if isinstance(image, Image.Image):
+                image.save(filepath, "PNG", quality=100)
+            else:
+                image = Image.fromarray(image)
+                image.save(filepath, "PNG", quality=100)
+            
+            if not os.path.exists(filepath):
+                print(f"Warning: Failed to verify saved image at {filepath}")
+                return None
+                
+            return filepath
+        except Exception as e:
+            print(f"Failed to save image: {str(e)}")
             return None
             
-        return filepath
     except Exception as e:
-        print(f"Error saving image: {str(e)}")
+        print(f"Error in save_image: {str(e)}")
         return None
 
+
 def load_gallery():
     """Load all images from the gallery directory"""
     try:
@@ -363,4 +370,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(allowed_paths=["/data"])  # /data 디렉토리 전체에 대한 접근 허용
+    demo.launch(allowed_paths=[PERSISTENT_DIR])