Update app.py

app.py

@@ -10,16 +10,20 @@ from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file, save_file
 import torch
 
-# --- Conversion Function: Safetensors → FP8 Safetensors (E4M3FN or E5M2) ---
+# 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()):
-    """
-    Loads a .safetensors file and saves a pruned FP8 version.
-    fp8_format: 'e4m3fn' or 'e5m2'
-    """
     progress(0.1, desc="Starting FP8 conversion...")
 
     try:
-        # Read metadata
         def read_safetensors_metadata(path):
             with open(path, 'rb') as f:
                 header_size = int.from_bytes(f.read(8), 'little')
@@ -30,28 +34,23 @@ def convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progres
         metadata = read_safetensors_metadata(safetensors_path)
         progress(0.3, desc="Loaded model metadata.")
 
-        # Load state dict
         state_dict = load_file(safetensors_path)
         progress(0.5, desc="Loaded model weights.")
 
-        # Select FP8 dtype
         if fp8_format == "e5m2":
             fp8_dtype = torch.float8_e5m2
-        else:  # default to e4m3fn
+        else:
             fp8_dtype = torch.float8_e4m3fn
 
-        # Convert to FP8
         sd_pruned = {}
         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})...")
-            # Only convert float tensors
             if state_dict[key].dtype in [torch.float16, torch.float32, torch.bfloat16]:
                 sd_pruned[key] = state_dict[key].to(fp8_dtype)
             else:
-                sd_pruned[key] = state_dict[key]  # keep non-float as-is (e.g., int for embeddings)
+                sd_pruned[key] = state_dict[key]
 
-        # Save FP8 safetensors
         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})
@@ -63,56 +62,165 @@ def convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progres
     except Exception as e:
         return False, str(e)
 
+# --- Source download helper ---
+def download_safetensors_file(
+    source_type,
+    repo_url,
+    filename,
+    hf_token=None,
+    modelscope_token=None,
+    progress=gr.Progress()
+):
+    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/", "")
+            safetensors_path = hf_hub_download(
+                repo_id=src_repo_id,
+                filename=filename,
+                cache_dir=temp_dir,
+                token=hf_token
+            )
+        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
+                )
+        else:
+            raise ValueError("Unknown source type")
+
+        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()
+):
+    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"
+        )
+        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)
+        # ModelScope requires model_type and license
+        api.push_model(
+            model_id=new_repo_id,
+            model_dir=output_dir,
+            commit_message=f"Upload FP8 ({fp8_format}) model"
+        )
+        return f"https://modelscope.cn/models/{new_repo_id}"
+
+    else:
+        raise ValueError("Unknown target type")
+
 # --- Main Processing Function ---
-def process_and_upload_fp8(repo_url, safetensors_filename, fp8_format, hf_token, new_repo_id, private_repo, progress=gr.Progress()):
-    if not all([repo_url, safetensors_filename, fp8_format, hf_token, new_repo_id]):
-        return None, "❌ Error: Please fill in all fields.", ""
+def process_and_upload_fp8(
+    source_type,
+    repo_url,
+    safetensors_filename,
+    fp8_format,
+    target_type,
+    new_repo_id,
+    hf_token,
+    modelscope_token,
+    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, "❌ Error: Invalid repository ID format. Use 'username/model-name'.", ""
+        return None, "❌ Invalid repository ID format. Use 'username/model-name'.", ""
 
-    temp_dir = tempfile.mkdtemp()
+    temp_dir = None
     output_dir = tempfile.mkdtemp()
 
     try:
-        # Authenticate
-        progress(0.05, desc="Logging into Hugging Face...")
-        api = HfApi(token=hf_token)
-        user_info = api.whoami()
-        user_name = user_info['name']
-        progress(0.1, desc=f"Logged in as {user_name}.")
-
-        # Parse source repo
-        clean_url = repo_url.strip().rstrip("/")
-        if "huggingface.co" not in clean_url:
-            return None, "❌ Source must be a Hugging Face model repo.", ""
-        src_repo_id = clean_url.replace("https://huggingface.co/", "")
-
-        # Download specified safetensors file
-        progress(0.15, desc=f"Downloading {safetensors_filename}...")
-        safetensors_path = hf_hub_download(
-            repo_id=src_repo_id,
+        # Authenticate & download
+        progress(0.05, desc="Authenticating and downloading...")
+        safetensors_path, temp_dir = download_safetensors_file(
+            source_type=source_type,
+            repo_url=repo_url,
             filename=safetensors_filename,
-            cache_dir=temp_dir,
-            token=hf_token
+            hf_token=hf_token,
+            modelscope_token=modelscope_token,
+            progress=progress
         )
         progress(0.25, desc="Download complete.")
 
-        # Convert to FP8
+        # Convert
         success, msg = convert_safetensors_to_fp8(safetensors_path, output_dir, fp8_format, progress)
         if not success:
             return None, f"❌ Conversion failed: {msg}", ""
 
-        # Create new repo
-        progress(0.92, desc="Creating new repository...")
-        api.create_repo(
-            repo_id=new_repo_id,
-            private=private_repo,
-            repo_type="model",
-            exist_ok=True
+        # Upload
+        progress(0.92, desc="Uploading model...")
+        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
         )
 
-        # Generate README
+        # README
         base_name = os.path.splitext(safetensors_filename)[0]
         fp8_filename = f"{base_name}-fp8-{fp8_format}.safetensors"
         readme = f"""---
@@ -128,77 +236,90 @@ tags:
 
 # FP8 Pruned Model ({fp8_format.upper()})
 
-Converted from: [`{src_repo_id}`](https://huggingface.co/{src_repo_id})  
+Converted from: `{repo_url}`  
 File: `{safetensors_filename}` → `{fp8_filename}`
 
 Quantization: **FP8 ({fp8_format.upper()})**  
-Converted by: {user_name}  
-Date: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+Converted on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
 
-> ⚠️ FP8 models require PyTorch ≥ 2.1 and compatible hardware (e.g., NVIDIA Ada/Hopper) for full acceleration. May fall back to FP16 on older GPUs.
+> ⚠️ Requires PyTorch ≥ 2.1 and compatible hardware for FP8 acceleration.
 """
-        with open(os.path.join(output_dir, "README.md"), "w") as f:
+        readme_path = os.path.join(output_dir, "README.md")
+        with open(readme_path, "w") as f:
             f.write(readme)
 
-        # Upload
-        progress(0.95, desc="Uploading to Hugging Face Hub...")
-        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}) pruned safetensors model"
-        )
+        # 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,
+                path_in_repo="README.md",
+                repo_id=new_repo_id,
+                repo_type="model",
+                token=hf_token
+            )
 
         progress(1.0, desc="✅ Done!")
         result_html = f"""
 ✅ Success!  
-Your FP8 ({fp8_format}) model is uploaded to: [{new_repo_id}](https://huggingface.co/{new_repo_id})  
-Visibility: {'Private' if private_repo else 'Public'}
+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()}
 """
         return gr.HTML(result_html), "✅ FP8 conversion and upload successful!", ""
 
     except Exception as e:
         return None, f"❌ Error: {str(e)}", ""
     finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
+        if temp_dir:
+            shutil.rmtree(temp_dir, ignore_errors=True)
         shutil.rmtree(output_dir, ignore_errors=True)
 
 # --- Gradio UI ---
-with gr.Blocks(title="Safetensors → FP8 Pruner") as demo:
+with gr.Blocks(title="Safetensors → FP8 Pruner (HF + ModelScope)") as demo:
     gr.Markdown("# 🔄 Safetensors to FP8 Pruner")
-    gr.Markdown("Converts any `.safetensors` file from a Hugging Face model repo to **FP8 (E4M3FN or E5M2)** for compact storage and faster inference.")
+    gr.Markdown("Convert `.safetensors` models to **FP8** and upload to **Hugging Face** or **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 Model Repository URL",
-                placeholder="https://huggingface.co/Yabo/FramePainter",
-                info="Hugging Face model repo containing your safetensors file"
+                label="Source Repository URL",
+                placeholder="e.g., https://huggingface.co/Yabo/FramePainter OR your-modelscope-id",
+                info="Hugging Face URL or ModelScope model ID"
             )
             safetensors_filename = gr.Textbox(
                 label="Safetensors Filename",
-                placeholder="unet_diffusion_pytorch_model.safetensors",
-                info="Name of the .safetensors file in the repo"
+                placeholder="unet_diffusion_pytorch_model.safetensors"
             )
             fp8_format = gr.Radio(
                 choices=["e4m3fn", "e5m2"],
                 value="e5m2",
                 label="FP8 Format",
-                info="E5M2 has wider dynamic range; E4M3FN has higher precision near zero."
+                info="E5M2: wider range; E4M3FN: better near-zero precision"
             )
             hf_token = gr.Textbox(
-                label="Hugging Face Token",
+                label="Hugging Face Token (if using HF)",
+                type="password"
+            )
+            modelscope_token = gr.Textbox(
+                label="ModelScope Token (optional)",
                 type="password",
-                info="Write-access token from https://huggingface.co/settings/tokens"
+                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",
-                info="Format: username/model-name"
+                placeholder="your-username/my-model-fp8"
             )
-            private_repo = gr.Checkbox(label="Make Private", value=False)
+            private_repo = gr.Checkbox(label="Make Private (HF only)", value=False)
 
     convert_btn = gr.Button("🚀 Convert & Upload", variant="primary")
 
@@ -208,16 +329,26 @@ with gr.Blocks(title="Safetensors → FP8 Pruner") as demo:
 
     convert_btn.click(
         fn=process_and_upload_fp8,
-        inputs=[repo_url, safetensors_filename, fp8_format, hf_token, new_repo_id, private_repo],
+        inputs=[
+            source_type,
+            repo_url,
+            safetensors_filename,
+            fp8_format,
+            target_type,
+            new_repo_id,
+            hf_token,
+            modelscope_token,
+            private_repo
+        ],
         outputs=[repo_link_output, status_output],
         show_progress=True
     )
 
     gr.Examples(
         examples=[
-            ["https://huggingface.co/Yabo/FramePainter", "unet_diffusion_pytorch_model.safetensors", "e5m2"]
+            ["huggingface", "https://huggingface.co/Yabo/FramePainter", "unet_diffusion_pytorch_model.safetensors", "e5m2", "huggingface"]
         ],
-        inputs=[repo_url, safetensors_filename, fp8_format]
+        inputs=[source_type, repo_url, safetensors_filename, fp8_format, target_type]
     )
 
 demo.launch()