Create app.py

app.py

@@ -0,0 +1,223 @@
+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
+
+# --- Conversion Function: Safetensors → FP8 Safetensors (E4M3FN or E5M2) ---
+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')
+                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.")
+
+        # 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
+            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)
+
+        # 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})
+        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)
+
+# --- 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.", ""
+
+    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'.", ""
+
+    temp_dir = tempfile.mkdtemp()
+    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,
+            filename=safetensors_filename,
+            cache_dir=temp_dir,
+            token=hf_token
+        )
+        progress(0.25, desc="Download complete.")
+
+        # Convert to FP8
+        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
+        )
+
+        # Generate 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: [`{src_repo_id}`](https://huggingface.co/{src_repo_id})  
+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')}
+
+> ⚠️ 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.
+"""
+        with open(os.path.join(output_dir, "README.md"), "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"
+        )
+
+        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'}
+"""
+        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)
+        shutil.rmtree(output_dir, ignore_errors=True)
+
+# --- Gradio UI ---
+with gr.Blocks(title="Safetensors → FP8 Pruner") 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.")
+
+    with gr.Row():
+        with gr.Column():
+            repo_url = gr.Textbox(
+                label="Source Model Repository URL",
+                placeholder="https://huggingface.co/Yabo/FramePainter",
+                info="Hugging Face model repo containing your safetensors file"
+            )
+            safetensors_filename = gr.Textbox(
+                label="Safetensors Filename",
+                placeholder="unet_diffusion_pytorch_model.safetensors",
+                info="Name of the .safetensors file in the repo"
+            )
+            fp8_format = gr.Radio(
+                choices=["e4m3fn", "e5m2"],
+                value="e5m2",
+                label="FP8 Format",
+                info="E5M2 has wider dynamic range; E4M3FN has higher precision near zero."
+            )
+            hf_token = gr.Textbox(
+                label="Hugging Face Token",
+                type="password",
+                info="Write-access token from https://huggingface.co/settings/tokens"
+            )
+        with gr.Column():
+            new_repo_id = gr.Textbox(
+                label="New Repository ID",
+                placeholder="your-username/my-model-fp8",
+                info="Format: username/model-name"
+            )
+            private_repo = gr.Checkbox(label="Make Private", 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=[repo_url, safetensors_filename, fp8_format, hf_token, new_repo_id, 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"]
+        ],
+        inputs=[repo_url, safetensors_filename, fp8_format]
+    )
+
+demo.launch()