Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import shutil
 from pathlib import Path
+import os
 from huggingface_hub import list_models
 
 # list_spaces is not guaranteed to exist in all HF runtimes
@@ -9,46 +10,89 @@ try:
 except ImportError:
     list_spaces = None
 
-
 # ===============================
-# Persistent storage
+# Configuration
 # ===============================
+OWNER_USERNAME = "your_hf_username"  # <-- replace with your username
+current_user = os.environ.get("HUGGINGFACE_HUB_USER", "")
+
 WORKFLOW_DIR = Path("/data/workflows")
-WORKFLOW_DIR.mkdir(parents=True, exist_ok=True)
+WORKFLOW_IMG_DIR = Path("/data/workflows_images")
+LORA_DIR = Path("/data/loras")
+LORA_IMG_DIR = Path("/data/loras_images")
+
+for d in [WORKFLOW_DIR, WORKFLOW_IMG_DIR, LORA_DIR, LORA_IMG_DIR]:
+    d.mkdir(parents=True, exist_ok=True)
 
-ALLOWED_EXTS = {".json", ".yaml", ".yml", ".png"}
+ALLOWED_FILE_EXTS = {".json", ".yaml", ".yml", ".png", ".safetensors"}
+ALLOWED_IMG_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
 
 
 # ===============================
-# Workflow upload handlers
+# Upload handlers
 # ===============================
-def upload_workflow(file):
+def upload_workflow(file, image=None):
     if file is None:
         return "No file uploaded."
-
     src = Path(file.name)
     ext = src.suffix.lower()
-
-    if ext not in ALLOWED_EXTS:
+    if ext not in ALLOWED_FILE_EXTS:
         return f"❌ Unsupported file type: `{ext}`"
-
     dest = WORKFLOW_DIR / src.name
     shutil.copy(src, dest)
-
+    # handle optional image
+    if image:
+        img_ext = Path(image.name).suffix.lower()
+        if img_ext in ALLOWED_IMG_EXTS:
+            shutil.copy(image.name, WORKFLOW_IMG_DIR / (src.stem + img_ext))
     return f"✅ Saved workflow: `{dest.name}`"
 
 
+def upload_lora(file, image=None):
+    if file is None:
+        return "No file uploaded."
+    src = Path(file.name)
+    ext = src.suffix.lower()
+    if ext not in ALLOWED_FILE_EXTS:
+        return f"❌ Unsupported file type: `{ext}`"
+    dest = LORA_DIR / src.name
+    shutil.copy(src, dest)
+    if image:
+        img_ext = Path(image.name).suffix.lower()
+        if img_ext in ALLOWED_IMG_EXTS:
+            shutil.copy(image.name, LORA_IMG_DIR / (src.stem + img_ext))
+    return f"✅ Saved LoRA: `{dest.name}`"
+
+
+# ===============================
+# Listing functions with images
+# ===============================
 def list_uploaded_workflows():
     files = sorted(WORKFLOW_DIR.glob("*"))
-
     if not files:
         return "No uploaded workflows yet."
+    md = ""
+    for f in files:
+        img_path = WORKFLOW_IMG_DIR / f"{f.stem}.png"
+        if img_path.exists():
+            md += f"### {f.name}\n![ref image]({img_path})\n---\n"
+        else:
+            md += f"### {f.name}\n(no image)\n---\n"
+    return md
 
-    return "\n---\n".join(
-        f"### {f.name}\n"
-        f"`{f.suffix}` | {f.stat().st_size // 1024} KB"
-        for f in files
-    )
+
+def list_uploaded_loras():
+    files = sorted(LORA_DIR.glob("*"))
+    if not files:
+        return "No uploaded LoRAs yet."
+    md = ""
+    for f in files:
+        img_path = LORA_IMG_DIR / f"{f.stem}.png"
+        if img_path.exists():
+            md += f"### {f.name}\n![ref image]({img_path})\n---\n"
+        else:
+            md += f"### {f.name}\n(no image)\n---\n"
+    return md
 
 
 # ===============================
@@ -61,10 +105,8 @@ def browse_loras(limit=20):
         direction=-1,
         limit=limit,
     ))
-
     if not models:
         return "No LoRAs found."
-
     return "\n---\n".join(
         f"### {m.modelId}\n"
         f"⬇️ {m.downloads or 0} | ⭐ {m.likes or 0}\n"
@@ -75,14 +117,12 @@ def browse_loras(limit=20):
 
 def browse_workflows(limit=20):
     uploaded_files = sorted(WORKFLOW_DIR.glob("*"))
-
     uploaded_md = ""
     if uploaded_files:
         uploaded_md = "## 🧑‍💻 Uploaded Workflows\n\n" + "\n---\n".join(
             f"### {f.name}\nLocal upload"
             for f in uploaded_files
         )
-
     if list_spaces is None:
         remote_md = "Spaces API unavailable."
     else:
@@ -91,25 +131,17 @@ def browse_workflows(limit=20):
             direction=-1,
             limit=limit,
         ))
-
         workflows = [
             s for s in spaces
-            if any(
-                tag in (s.tags or [])
-                for tag in ["comfyui", "workflow", "pipeline", "diffusers"]
-            )
+            if any(tag in (s.tags or []) for tag in ["comfyui", "workflow", "pipeline", "diffusers"])
         ]
-
         if workflows:
             remote_md = "## 🌍 Community Workflows\n\n" + "\n---\n".join(
-                f"### {s.id}\n"
-                f"⭐ {s.likes or 0}\n"
-                f"https://huggingface.co/spaces/{s.id}"
+                f"### {s.id}\n⭐ {s.likes or 0}\nhttps://huggingface.co/spaces/{s.id}"
                 for s in workflows[:limit]
             )
         else:
             remote_md = "No community workflows found."
-
     return uploaded_md + "\n\n" + remote_md
 
 
@@ -136,26 +168,40 @@ with gr.Blocks() as demo:
                 outputs=lora_out,
             )
 
-        # Upload Tab
-        with gr.Tab("⬆️ Upload Workflow"):
-            gr.Markdown(
-                "Upload workflow files (`.json`, `.yaml`, `.png`). "
-                "Files are saved persistently."
-            )
-
-            wf_file = gr.File(label="Workflow file")
-            upload_status = gr.Markdown()
-
-            gr.Button("Upload").click(
-                upload_workflow,
-                inputs=wf_file,
-                outputs=upload_status,
-            )
-
-            uploaded_list = gr.Markdown()
-            gr.Button("Refresh Uploaded Workflows").click(
-                list_uploaded_workflows,
-                outputs=uploaded_list,
-            )
+        # Upload tabs visible only to owner
+        if current_user == OWNER_USERNAME:
+            # Upload Workflow
+            with gr.Tab("⬆️ Upload Workflow"):
+                gr.Markdown("Upload workflow files with optional reference image (.json/.yaml/.png).")
+                wf_file = gr.File(label="Workflow file")
+                wf_image = gr.File(label="Reference image (optional)")
+                wf_status = gr.Markdown()
+                gr.Button("Upload").click(
+                    upload_workflow,
+                    inputs=[wf_file, wf_image],
+                    outputs=wf_status,
+                )
+                wf_list = gr.Markdown()
+                gr.Button("Refresh Uploaded Workflows").click(
+                    list_uploaded_workflows,
+                    outputs=wf_list,
+                )
+
+            # Upload LoRA
+            with gr.Tab("⬆️ Upload LoRA"):
+                gr.Markdown("Upload LoRA files with optional reference image (.safetensors/.ckpt/.png).")
+                lora_file = gr.File(label="LoRA file")
+                lora_image = gr.File(label="Reference image (optional)")
+                lora_status = gr.Markdown()
+                gr.Button("Upload").click(
+                    upload_lora,
+                    inputs=[lora_file, lora_image],
+                    outputs=lora_status,
+                )
+                lora_list = gr.Markdown()
+                gr.Button("Refresh Uploaded LoRAs").click(
+                    list_uploaded_loras,
+                    outputs=lora_list,
+                )
 
 demo.launch()