Update app.py

app.py

@@ -62,6 +62,37 @@ def convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progres
     except Exception as e:
         return False, str(e)
 
+# --- Parse HF URL with optional subfolder ---
+def parse_hf_url(url):
+    """
+    Parses a Hugging Face URL like:
+    - https://huggingface.co/username/repo
+    - https://huggingface.co/username/repo/tree/main/subfolder
+    Returns (repo_id, subfolder)
+    """
+    url = url.strip().rstrip("/")
+    if not url.startswith("https://huggingface.co/"):
+        raise ValueError("URL must start with https://huggingface.co/")
+    path = url.replace("https://huggingface.co/", "")
+    parts = path.split("/")
+
+    if len(parts) < 2:
+        raise ValueError("Invalid repo format")
+
+    # repo_id is always first two parts
+    repo_id = "/".join(parts[:2])
+
+    # Check if "/tree/branch/" is present
+    subfolder = ""
+    if len(parts) > 3 and parts[2] == "tree":
+        # everything after branch is subfolder
+        subfolder = "/".join(parts[4:]) if len(parts) > 4 else ""
+    elif len(parts) > 2:
+        # old style: username/repo/subfolder (not standard, but support)
+        subfolder = "/".join(parts[2:])
+
+    return repo_id, subfolder
+
 # --- Source download helper ---
 def download_safetensors_file(
     source_type,
@@ -74,15 +105,14 @@ def download_safetensors_file(
     temp_dir = tempfile.mkdtemp()
     try:
         if source_type == "huggingface":
-            clean_url = repo_url.strip().rstrip("/")
-            if "huggingface.co" not in clean_url:
-                raise ValueError("Invalid Hugging Face URL")
-            src_repo_id = clean_url.replace("https://huggingface.co/", "")
+            repo_id, subfolder = parse_hf_url(repo_url)
             safetensors_path = hf_hub_download(
-                repo_id=src_repo_id,
+                repo_id=repo_id,
                 filename=filename,
+                subfolder=subfolder or None,
                 cache_dir=temp_dir,
-                token=hf_token
+                token=hf_token,
+                resume_download=True
             )
         elif source_type == "modelscope":
             if not MODELScope_AVAILABLE:
@@ -148,7 +178,6 @@ def upload_to_target(
         api = ModelScopeApi()
         if modelscope_token:
             api.login(modelscope_token)
-        # ModelScope requires model_type and license
         api.push_model(
             model_id=new_repo_id,
             model_dir=output_dir,
@@ -190,8 +219,7 @@ def process_and_upload_fp8(
     output_dir = tempfile.mkdtemp()
 
     try:
-        # Authenticate & download
-        progress(0.05, desc="Authenticating and downloading...")
+        progress(0.05, desc="Parsing URL and downloading...")
         safetensors_path, temp_dir = download_safetensors_file(
             source_type=source_type,
             repo_url=repo_url,
@@ -202,12 +230,10 @@ def process_and_upload_fp8(
         )
         progress(0.25, desc="Download complete.")
 
-        # Convert
         success, msg = convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progress)
         if not success:
             return None, f"❌ Conversion failed: {msg}", ""
 
-        # Upload
         progress(0.92, desc="Uploading model...")
         repo_url_final = upload_to_target(
             target_type=target_type,
@@ -220,7 +246,6 @@ def process_and_upload_fp8(
             progress=progress
         )
 
-        # README
         base_name = os.path.splitext(safetensors_filename)[0]
         fp8_filename = f"{base_name}-fp8-{fp8_format}.safetensors"
         readme = f"""---
@@ -241,14 +266,11 @@ File: `{safetensors_filename}` → `{fp8_filename}`
 
 Quantization: **FP8 ({fp8_format.upper()})**  
 Converted on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
-
-> ⚠️ Requires PyTorch ≥ 2.1 and compatible hardware for FP8 acceleration.
 """
         readme_path = os.path.join(output_dir, "README.md")
         with open(readme_path, "w") as f:
             f.write(readme)
 
-        # Re-upload README if needed (for ModelScope, already included; for HF, upload separately)
         if target_type == "huggingface":
             HfApi(token=hf_token).upload_file(
                 path_or_fileobj=readme_path,
@@ -277,6 +299,7 @@ Source: {source_type.title()} → Target: {target_type.title()}
 with gr.Blocks(title="Safetensors → FP8 Pruner (HF + ModelScope)") as demo:
     gr.Markdown("# 🔄 Safetensors to FP8 Pruner")
     gr.Markdown("Convert `.safetensors` models to **FP8** and upload to **Hugging Face** or **ModelScope**.")
+    gr.Markdown("Supports subfolders: e.g., `https://huggingface.co/lixiaowen/diffuEraser/tree/main/brushnet`")
 
     with gr.Row():
         with gr.Column():
@@ -287,18 +310,17 @@ with gr.Blocks(title="Safetensors → FP8 Pruner (HF + ModelScope)") as demo:
             )
             repo_url = gr.Textbox(
                 label="Source Repository URL",
-                placeholder="e.g., https://huggingface.co/Yabo/FramePainter OR your-modelscope-id",
-                info="Hugging Face URL or ModelScope model ID"
+                placeholder="https://huggingface.co/lixiaowen/diffuEraser/tree/main/brushnet",
+                info="Full URL including subfolder (if any)"
             )
             safetensors_filename = gr.Textbox(
                 label="Safetensors Filename",
-                placeholder="unet_diffusion_pytorch_model.safetensors"
+                placeholder="diffusion_pytorch_model.safetensors"
             )
             fp8_format = gr.Radio(
                 choices=["e4m3fn", "e5m2"],
                 value="e5m2",
-                label="FP8 Format",
-                info="E5M2: wider range; E4M3FN: better near-zero precision"
+                label="FP8 Format"
             )
             hf_token = gr.Textbox(
                 label="Hugging Face Token (if using HF)",
@@ -346,7 +368,7 @@ with gr.Blocks(title="Safetensors → FP8 Pruner (HF + ModelScope)") as demo:
 
     gr.Examples(
         examples=[
-            ["huggingface", "https://huggingface.co/Yabo/FramePainter", "unet_diffusion_pytorch_model.safetensors", "e5m2", "huggingface"]
+            ["huggingface", "https://huggingface.co/lixiaowen/diffuEraser/tree/main/brushnet", "diffusion_pytorch_model.safetensors", "e5m2", "huggingface"]
         ],
         inputs=[source_type, repo_url, safetensors_filename, fp8_format, target_type]
     )