Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import tempfile | |
| import shutil | |
| import re | |
| import json | |
| import datetime | |
| from pathlib import Path | |
| 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...") | |
| try: | |
| def read_safetensors_metadata(path): | |
| with open(path, 'rb') as f: | |
| header_size = int.from_bytes(f.read(8), 'little') | |
| header_json = f.read(header_size).decode('utf-8') | |
| header = json.loads(header_json) | |
| return header.get('__metadata__', {}) | |
| metadata = read_safetensors_metadata(safetensors_path) | |
| progress(0.3, desc="Loaded model metadata.") | |
| state_dict = load_file(safetensors_path) | |
| progress(0.5, desc="Loaded model weights.") | |
| if fp8_format == "e5m2": | |
| fp8_dtype = torch.float8_e5m2 | |
| else: | |
| fp8_dtype = torch.float8_e4m3fn | |
| 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})...") | |
| 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] | |
| 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.") | |
| progress(1.0, desc="FP8 conversion complete!") | |
| return True, f"Model successfully pruned to FP8 ({fp8_format})." | |
| 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( | |
| 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, "❌ Invalid repository ID format. Use 'username/model-name'.", "" | |
| temp_dir = None | |
| output_dir = tempfile.mkdtemp() | |
| try: | |
| # 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, | |
| hf_token=hf_token, | |
| modelscope_token=modelscope_token, | |
| progress=progress | |
| ) | |
| 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, | |
| 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 | |
| ) | |
| # README | |
| 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 | |
| - diffusion | |
| - converted-by-gradio | |
| - fp8-{fp8_format} | |
| --- | |
| # FP8 Pruned Model ({fp8_format.upper()}) | |
| Converted from: `{repo_url}` | |
| 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, | |
| 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 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: | |
| 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 (HF + ModelScope)") as demo: | |
| gr.Markdown("# 🔄 Safetensors to FP8 Pruner") | |
| 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 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" | |
| ) | |
| fp8_format = gr.Radio( | |
| choices=["e4m3fn", "e5m2"], | |
| value="e5m2", | |
| label="FP8 Format", | |
| info="E5M2: wider range; E4M3FN: better near-zero precision" | |
| ) | |
| 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 | |
| ) | |
| 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) | |
| convert_btn = gr.Button("🚀 Convert & Upload", variant="primary") | |
| with gr.Row(): | |
| status_output = gr.Markdown() | |
| repo_link_output = gr.HTML() | |
| convert_btn.click( | |
| fn=process_and_upload_fp8, | |
| 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=[ | |
| ["huggingface", "https://huggingface.co/Yabo/FramePainter", "unet_diffusion_pytorch_model.safetensors", "e5m2", "huggingface"] | |
| ], | |
| inputs=[source_type, repo_url, safetensors_filename, fp8_format, target_type] | |
| ) | |
| demo.launch() |