Update app.py

app.py

@@ -10,18 +10,15 @@ from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file, save_file
 import torch
 
-# Optional ModelScope integration
 try:
-    from modelscope.hub.snapshot_download import snapshot_download as ms_snapshot_download
     from modelscope.hub.file_download import model_file_download as ms_file_download
     from modelscope.hub.api import HubApi as ModelScopeApi
     MODELScope_AVAILABLE = True
 except ImportError:
     MODELScope_AVAILABLE = False
 
-# --- Conversion Function: Safetensors → FP8 Safetensors ---
-def convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progress=gr.Progress()):
-    progress(0.1, desc="Starting FP8 conversion...")
+def convert_safetensors_to_fp8_with_delta(safetensors_path, output_dir, fp8_format, progress=gr.Progress()):
+    progress(0.1, desc="Starting FP8 conversion with delta...")
 
     try:
         def read_safetensors_metadata(path):
@@ -32,76 +29,63 @@ def convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progres
                 return header.get('__metadata__', {})
 
         metadata = read_safetensors_metadata(safetensors_path)
-        progress(0.3, desc="Loaded model metadata.")
+        progress(0.2, desc="Loaded metadata.")
 
         state_dict = load_file(safetensors_path)
-        progress(0.5, desc="Loaded model weights.")
+        progress(0.4, desc="Loaded weights.")
 
         if fp8_format == "e5m2":
             fp8_dtype = torch.float8_e5m2
         else:
             fp8_dtype = torch.float8_e4m3fn
 
-        sd_pruned = {}
+        sd_fp8 = {}
+        sd_delta = {}
         total = len(state_dict)
+
         for i, key in enumerate(state_dict):
-            progress(0.5 + 0.4 * (i / total), desc=f"Converting tensor {i+1}/{total} to FP8 ({fp8_format})...")
-            if state_dict[key].dtype in [torch.float16, torch.float32, torch.bfloat16]:
-                sd_pruned[key] = state_dict[key].to(fp8_dtype)
+            progress(0.4 + 0.4 * (i / total), desc=f"Processing {i+1}/{total}...")
+            weight = state_dict[key]
+            if weight.dtype in [torch.float16, torch.float32, torch.bfloat16]:
+                fp8_weight = weight.to(fp8_dtype)
+                fp8_recon = fp8_weight.to(weight.dtype)
+                delta = weight - fp8_recon
+                sd_fp8[key] = fp8_weight
+                sd_delta[f"delta.{key}"] = delta
             else:
-                sd_pruned[key] = state_dict[key]
+                sd_fp8[key] = weight
 
         base_name = os.path.splitext(os.path.basename(safetensors_path))[0]
-        output_path = os.path.join(output_dir, f"{base_name}-fp8-{fp8_format}.safetensors")
-        save_file(sd_pruned, output_path, metadata={"format": "pt", "fp8_format": fp8_format, **metadata})
-        progress(0.9, desc="Saved FP8 safetensors file.")
+        fp8_path = os.path.join(output_dir, f"{base_name}-fp8-{fp8_format}.safetensors")
+        delta_path = os.path.join(output_dir, f"{base_name}-fp8-delta.safetensors")
+
+        save_file(sd_fp8, fp8_path, metadata={"format": "pt", "fp8_format": fp8_format, **metadata})
+        save_file(sd_delta, delta_path, metadata={"format": "pt", "source": "fp8_delta", "fp8_format": fp8_format})
 
-        progress(1.0, desc="FP8 conversion complete!")
-        return True, f"Model successfully pruned to FP8 ({fp8_format})."
+        progress(0.9, desc="Saved FP8 and delta files.")
+        progress(1.0, desc="✅ FP8 + delta generation complete!")
+        return True, f"FP8 ({fp8_format}) and delta saved."
 
     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,
-    repo_url,
-    filename,
-    hf_token=None,
-    modelscope_token=None,
-    progress=gr.Progress()
-):
+def download_safetensors_file(source_type, repo_url, filename, hf_token=None, progress=gr.Progress()):
     temp_dir = tempfile.mkdtemp()
     try:
         if source_type == "huggingface":
@@ -116,79 +100,31 @@ def download_safetensors_file(
             )
         elif source_type == "modelscope":
             if not MODELScope_AVAILABLE:
-                raise ImportError("ModelScope not installed. Install with: pip install modelscope")
-            clean_url = repo_url.strip().rstrip("/")
-            if "modelscope.cn" in clean_url:
-                src_repo_id = "/".join(clean_url.split("/")[-2:])
-            else:
-                src_repo_id = repo_url.strip()
-            if modelscope_token:
-                os.environ["MODELSCOPE_CACHE"] = temp_dir
-                safetensors_path = ms_file_download(
-                    model_id=src_repo_id,
-                    file_path=filename,
-                    token=modelscope_token
-                )
-            else:
-                safetensors_path = ms_file_download(
-                    model_id=src_repo_id,
-                    file_path=filename
-                )
+                raise ImportError("ModelScope not installed")
+            repo_id = repo_url.strip()
+            safetensors_path = ms_file_download(model_id=repo_id, file_path=filename)
         else:
-            raise ValueError("Unknown source type")
-
+            raise ValueError("Unknown source")
         return safetensors_path, temp_dir
     except Exception as e:
         shutil.rmtree(temp_dir, ignore_errors=True)
         raise e
 
-# --- Upload helper ---
-def upload_to_target(
-    target_type,
-    new_repo_id,
-    output_dir,
-    fp8_format,
-    hf_token=None,
-    modelscope_token=None,
-    private_repo=False,
-    progress=gr.Progress()
-):
+def upload_to_target(target_type, new_repo_id, output_dir, fp8_format, hf_token=None, modelscope_token=None, private_repo=False):
     if target_type == "huggingface":
-        if not hf_token:
-            raise ValueError("Hugging Face token required")
         api = HfApi(token=hf_token)
-        api.create_repo(
-            repo_id=new_repo_id,
-            private=private_repo,
-            repo_type="model",
-            exist_ok=True
-        )
-        api.upload_folder(
-            repo_id=new_repo_id,
-            folder_path=output_dir,
-            repo_type="model",
-            token=hf_token,
-            commit_message=f"Upload FP8 ({fp8_format}) model"
-        )
+        api.create_repo(repo_id=new_repo_id, private=private_repo, repo_type="model", exist_ok=True)
+        api.upload_folder(repo_id=new_repo_id, folder_path=output_dir, repo_type="model", token=hf_token)
         return f"https://huggingface.co/{new_repo_id}"
-
     elif target_type == "modelscope":
-        if not MODELScope_AVAILABLE:
-            raise ImportError("ModelScope not installed")
         api = ModelScopeApi()
         if modelscope_token:
             api.login(modelscope_token)
-        api.push_model(
-            model_id=new_repo_id,
-            model_dir=output_dir,
-            commit_message=f"Upload FP8 ({fp8_format}) model"
-        )
+        api.push_model(model_id=new_repo_id, model_dir=output_dir)
         return f"https://modelscope.cn/models/{new_repo_id}"
-
     else:
-        raise ValueError("Unknown target type")
+        raise ValueError("Unknown target")
 
-# --- Main Processing Function ---
 def process_and_upload_fp8(
     source_type,
     repo_url,
@@ -201,79 +137,80 @@ def process_and_upload_fp8(
     private_repo,
     progress=gr.Progress()
 ):
-    required_fields = [repo_url, safetensors_filename, new_repo_id]
-    if source_type == "huggingface":
-        required_fields.append(hf_token)
-    if target_type == "huggingface":
-        required_fields.append(hf_token)
-    if target_type == "modelscope" and modelscope_token:
-        required_fields.append(modelscope_token)
-
-    if not all(required_fields):
-        return None, "❌ Error: Please fill in all required fields.", ""
-
     if not re.match(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$", new_repo_id):
-        return None, "❌ Invalid repository ID format. Use 'username/model-name'.", ""
+        return None, "❌ Invalid repo ID format. Use 'username/model-name'.", ""
+
+    if source_type == "huggingface" and not hf_token:
+        return None, "❌ Hugging Face token required for source.", ""
+    if target_type == "huggingface" and not hf_token:
+        return None, "❌ Hugging Face token required for target.", ""
 
     temp_dir = None
     output_dir = tempfile.mkdtemp()
 
     try:
-        progress(0.05, desc="Parsing URL and downloading...")
+        progress(0.05, desc="Downloading model...")
         safetensors_path, temp_dir = download_safetensors_file(
-            source_type=source_type,
-            repo_url=repo_url,
-            filename=safetensors_filename,
-            hf_token=hf_token,
-            modelscope_token=modelscope_token,
-            progress=progress
+            source_type, repo_url, safetensors_filename, hf_token, progress
         )
-        progress(0.25, desc="Download complete.")
 
-        success, msg = convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progress)
+        progress(0.25, desc="Converting to FP8 with delta...")
+        success, msg = convert_safetensors_to_fp8_with_delta(safetensors_path, output_dir, fp8_format, progress)
         if not success:
             return None, f"❌ Conversion failed: {msg}", ""
 
-        progress(0.92, desc="Uploading model...")
+        progress(0.9, desc="Uploading...")
         repo_url_final = upload_to_target(
-            target_type=target_type,
-            new_repo_id=new_repo_id,
-            output_dir=output_dir,
-            fp8_format=fp8_format,
-            hf_token=hf_token,
-            modelscope_token=modelscope_token,
-            private_repo=private_repo,
-            progress=progress
+            target_type, new_repo_id, output_dir, fp8_format, hf_token, modelscope_token, private_repo
         )
 
         base_name = os.path.splitext(safetensors_filename)[0]
-        fp8_filename = f"{base_name}-fp8-{fp8_format}.safetensors"
         readme = f"""---
 library_name: diffusers
 tags:
 - fp8
 - safetensors
-- pruned
+- delta-compensation
 - diffusion
 - converted-by-gradio
-- fp8-{fp8_format}
 ---
 
-# FP8 Pruned Model ({fp8_format.upper()})
+# FP8 Model with Delta Compensation
+
+- **Source**: `{repo_url}`
+- **File**: `{safetensors_filename}`
+- **FP8 Format**: `{fp8_format.upper()}`
+- **Delta File**: `{base_name}-fp8-delta.safetensors`
+
+## Usage (Inference)
+
+To restore near-original precision:
+
+```python
+import torch
+from safetensors.torch import load_file
+
+fp8_state = load_file("{base_name}-fp8-{fp8_format}.safetensors")
+delta_state = load_file("{base_name}-fp8-delta.safetensors")
 
-Converted from: `{repo_url}`  
-File: `{safetensors_filename}` → `{fp8_filename}`
+restored_state = {{}}
+for key in fp8_state:
+    if f"delta.{{key}}" in delta_state:
+        fp8_weight = fp8_state[key].to(torch.float32)
+        delta = delta_state[f"delta.{{key}}"]
+        restored_state[key] = fp8_weight + delta
+    else:
+        restored_state[key] = fp8_state[key].to(torch.float32)
+```
 
-Quantization: **FP8 ({fp8_format.upper()})**  
-Converted on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+> Requires PyTorch ≥ 2.1 for FP8 support.
 """
-        readme_path = os.path.join(output_dir, "README.md")
-        with open(readme_path, "w") as f:
+        with open(os.path.join(output_dir, "README.md"), "w") as f:
             f.write(readme)
 
         if target_type == "huggingface":
             HfApi(token=hf_token).upload_file(
-                path_or_fileobj=readme_path,
+                path_or_fileobj=os.path.join(output_dir, "README.md"),
                 path_in_repo="README.md",
                 repo_id=new_repo_id,
                 repo_type="model",
@@ -283,10 +220,10 @@ Converted on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
         progress(1.0, desc="✅ Done!")
         result_html = f"""
 ✅ Success!  
-Your FP8 model is uploaded to: <a href="{repo_url_final}" target="_blank">{new_repo_id}</a>  
-Source: {source_type.title()} → Target: {target_type.title()}
+Model uploaded to: <a href="{repo_url_final}" target="_blank">{new_repo_id}</a>  
+Includes: FP8 model + delta compensation file.
 """
-        return gr.HTML(result_html), "✅ FP8 conversion and upload successful!", ""
+        return gr.HTML(result_html), "✅ FP8 + delta upload successful!", ""
 
     except Exception as e:
         return None, f"❌ Error: {str(e)}", ""
@@ -295,59 +232,26 @@ Source: {source_type.title()} → Target: {target_type.title()}
             shutil.rmtree(temp_dir, ignore_errors=True)
         shutil.rmtree(output_dir, ignore_errors=True)
 
-# --- Gradio UI ---
-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.Blocks(title="FP8 + Delta Converter (HF ↔ ModelScope)") as demo:
+    gr.Markdown("# 🔄 FP8 Pruner with Delta Compensation")
+    gr.Markdown("Convert `.safetensors` → **FP8** + **delta file** for precision recovery. Supports Hugging Face ↔ ModelScope.")
 
     with gr.Row():
         with gr.Column():
-            source_type = gr.Radio(
-                choices=["huggingface", "modelscope"],
-                value="huggingface",
-                label="Source Platform"
-            )
-            repo_url = gr.Textbox(
-                label="Source Repository URL",
-                placeholder="https://huggingface.co/lixiaowen/diffuEraser/tree/main/brushnet",
-                info="Full URL including subfolder (if any)"
-            )
-            safetensors_filename = gr.Textbox(
-                label="Safetensors Filename",
-                placeholder="diffusion_pytorch_model.safetensors"
-            )
-            fp8_format = gr.Radio(
-                choices=["e4m3fn", "e5m2"],
-                value="e5m2",
-                label="FP8 Format"
-            )
-            hf_token = gr.Textbox(
-                label="Hugging Face Token (if using HF)",
-                type="password"
-            )
-            modelscope_token = gr.Textbox(
-                label="ModelScope Token (optional)",
-                type="password",
-                visible=MODELScope_AVAILABLE
-            )
+            source_type = gr.Radio(["huggingface", "modelscope"], value="huggingface", label="Source")
+            repo_url = gr.Textbox(label="Repo URL or ID", placeholder="https://huggingface.co/... or modelscope-id")
+            safetensors_filename = gr.Textbox(label="Filename", placeholder="model.safetensors")
+            fp8_format = gr.Radio(["e4m3fn", "e5m2"], value="e5m2", label="FP8 Format")
+            hf_token = gr.Textbox(label="HF Token (only if using HF)", type="password")
+            modelscope_token = gr.Textbox(label="ModelScope Token (optional)", type="password", visible=MODELScope_AVAILABLE)
         with gr.Column():
-            target_type = gr.Radio(
-                choices=["huggingface", "modelscope"],
-                value="huggingface",
-                label="Target Platform"
-            )
-            new_repo_id = gr.Textbox(
-                label="New Repository ID",
-                placeholder="your-username/my-model-fp8"
-            )
-            private_repo = gr.Checkbox(label="Make Private (HF only)", value=False)
+            target_type = gr.Radio(["huggingface", "modelscope"], value="huggingface", label="Target")
+            new_repo_id = gr.Textbox(label="New Repo ID", placeholder="user/model-fp8")
+            private_repo = gr.Checkbox(label="Private (HF only)", value=False)
 
     convert_btn = gr.Button("🚀 Convert & Upload", variant="primary")
-
-    with gr.Row():
-        status_output = gr.Markdown()
-        repo_link_output = gr.HTML()
+    status_output = gr.Markdown()
+    repo_link_output = gr.HTML()
 
     convert_btn.click(
         fn=process_and_upload_fp8,
@@ -368,7 +272,7 @@ with gr.Blocks(title="Safetensors → FP8 Pruner (HF + ModelScope)") as demo:
 
     gr.Examples(
         examples=[
-            ["huggingface", "https://huggingface.co/lixiaowen/diffuEraser/tree/main/brushnet", "diffusion_pytorch_model.safetensors", "e5m2", "huggingface"]
+            ["huggingface", "https://huggingface.co/Yabo/FramePainter/tree/main", "unet_diffusion_pytorch_model.safetensors", "e5m2", "modelscope"]
         ],
         inputs=[source_type, repo_url, safetensors_filename, fp8_format, target_type]
     )