Update app.py

app.py

@@ -1,56 +1,100 @@
+import os
+import time
+import uuid
+from typing import Optional
+
 import gradio as gr
 from PIL import Image
-import os, time, uuid
 
-# If your Space has "Persistent Storage" enabled, /data is the mount point.
-# If not, it still works but files won't survive a restart.
-SAVE_DIR = os.getenv("SAVE_DIR", "/data/uploads")
-os.makedirs(SAVE_DIR, exist_ok=True)
+# ------------------------------------------------------------
+# Pick a writable save directory automatically.
+# Order:
+# 1) SAVE_DIR env var (if supplied)
+# 2) /data/uploads           (when Persistent Storage is enabled)
+# 3) /home/user/app/uploads  (always writable in Spaces)
+# 4) /tmp/uploads            (always writable, ephemeral)
+# ------------------------------------------------------------
+def pick_writable_dir() -> str:
+    candidates = []
+    env_dir = os.getenv("SAVE_DIR")
+    if env_dir:
+        candidates.append(env_dir)
+    candidates += ["/data/uploads", "/home/user/app/uploads", "/tmp/uploads"]
+
+    for path in candidates:
+        try:
+            os.makedirs(path, exist_ok=True)
+            # double-check we can create files here
+            testfile = os.path.join(path, f".write_test_{uuid.uuid4().hex}")
+            with open(testfile, "w") as f:
+                f.write("ok")
+            os.remove(testfile)
+            return path
+        except Exception:
+            continue
+    # last resort: current directory
+    fallback = "./uploads"
+    os.makedirs(fallback, exist_ok=True)
+    return fallback
 
-ALLOWED_FORMATS = {"JPEG", "PNG"}   # strict enforcement
+SAVE_DIR = pick_writable_dir()
 
-def handle_upload(file_path: str):
+ALLOWED_FORMATS = {"JPEG", "PNG"}  # strict type enforcement
+
+def handle_upload(file_path: Optional[str]):
     if not file_path:
         return "Please choose an image.", None
 
-    # Verify the real file type by reading the header via Pillow.
     try:
         with Image.open(file_path) as im:
             fmt = (im.format or "").upper()
             if fmt not in ALLOWED_FORMATS:
                 return f"Only JPEG/PNG allowed. Detected: {fmt or 'unknown'}.", None
 
-            # Optional: normalize mode (e.g., ensure RGB)
+            # Normalize to RGB so we can always save safely
             if im.mode not in ("RGB", "RGBA"):
                 im = im.convert("RGB")
 
-            # Save with a unique filename
+            # Choose extension from the detected format
             ext = ".jpg" if fmt == "JPEG" else ".png"
             filename = f"{int(time.time())}_{uuid.uuid4().hex}{ext}"
             save_path = os.path.join(SAVE_DIR, filename)
-            im.save(save_path)
 
-            info = f"✅ Saved as **{filename}** · {im.width}×{im.height}px · {fmt}"
-            # Return a copy for the preview (so file handle is closed)
+            # Reasonable defaults for each format
+            if fmt == "JPEG":
+                im.save(save_path, quality=90, optimize=True)
+            else:
+                im.save(save_path, optimize=True)
+
+            info = (
+                f"✅ Saved to **{save_path}**\n\n"
+                f"- File: **{filename}**\n"
+                f"- Size: **{im.width}×{im.height}px**\n"
+                f"- Format: **{fmt}**"
+            )
             return info, im.copy()
     except Exception as e:
-        return f"Failed to read image: {e}", None
+        return f"Failed to read/save image: {e}", None
 
 with gr.Blocks(title="Image Upload (JPEG/PNG)") as demo:
-    gr.Markdown("### Upload an image (JPEG/PNG). It will be validated and saved to storage.")
+    gr.Markdown("### Upload a JPEG or PNG. The file will be validated and saved.")
     with gr.Row():
         inp = gr.Image(
             label="Drop or click to upload",
-            type="filepath",               # gives us the original temp filepath
+            type="filepath",              # keep original temp filepath
             image_mode="RGB",
             sources=["upload", "clipboard"]
         )
         preview = gr.Image(label="Preview", interactive=False)
 
     msg = gr.Markdown()
-    # Auto-save on change; also provide an explicit button for clarity.
+
+    # Auto-save when the image changes; also allow manual click.
     inp.change(handle_upload, inputs=inp, outputs=[msg, preview])
     gr.Button("Save").click(handle_upload, inputs=inp, outputs=[msg, preview])
 
+    gr.Markdown(f"**Save directory:** `{SAVE_DIR}`")
+
 if __name__ == "__main__":
     demo.launch()
+