Update app.py

app.py

@@ -9,6 +9,7 @@ from pathlib import Path
 from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file, save_file
 import torch
+import torch.nn.functional as F
 
 try:
     from modelscope.hub.file_download import model_file_download as ms_file_download
@@ -17,8 +18,16 @@ try:
 except ImportError:
     MODELScope_AVAILABLE = False
 
-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...")
+def low_rank_decomposition(weight, rank=64):
+    if weight.ndim != 2:
+        return None
+    U, S, Vh = torch.linalg.svd(weight.float(), full_matrices=False)
+    U = U[:, :rank] @ torch.diag(torch.sqrt(S[:rank]))
+    Vh = torch.diag(torch.sqrt(S[:rank])) @ Vh[:rank, :]
+    return U.contiguous(), Vh.contiguous()
+
+def convert_safetensors_to_fp8_with_lora(safetensors_path, output_dir, fp8_format, lora_rank=64, progress=gr.Progress()):
+    progress(0.1, desc="Starting FP8 conversion with LoRA extraction...")
 
     try:
         def read_safetensors_metadata(path):
@@ -40,31 +49,41 @@ def convert_safetensors_to_fp8_with_delta(safetensors_path, output_dir, fp8_form
             fp8_dtype = torch.float8_e4m3fn
 
         sd_fp8 = {}
-        sd_delta = {}
+        lora_weights = {}
         total = len(state_dict)
+        lora_keys = []
 
         for i, key in enumerate(state_dict):
             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
+
+                # Attempt LoRA decomposition only for 2D tensors
+                if weight.ndim == 2 and min(weight.shape) > lora_rank:
+                    try:
+                        U, V = low_rank_decomposition(weight, rank=lora_rank)
+                        if U is not None and V is not None:
+                            lora_weights[f"lora_A.{key}"] = U.to(torch.float16)
+                            lora_weights[f"lora_B.{key}"] = V.to(torch.float16)
+                            lora_keys.append(key)
+                    except Exception:
+                        pass
             else:
                 sd_fp8[key] = weight
 
         base_name = os.path.splitext(os.path.basename(safetensors_path))[0]
         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")
+        lora_path = os.path.join(output_dir, f"{base_name}-lora-r{lora_rank}.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})
+        if lora_weights:
+            save_file(lora_weights, lora_path, metadata={"format": "pt", "lora_rank": str(lora_rank)})
 
-        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."
+        progress(0.9, desc="Saved FP8 and LoRA files.")
+        progress(1.0, desc="✅ FP8 + LoRA extraction complete!")
+        return True, f"FP8 ({fp8_format}) and rank-{lora_rank} LoRA saved."
 
     except Exception as e:
         return False, str(e)
@@ -130,6 +149,7 @@ def process_and_upload_fp8(
     repo_url,
     safetensors_filename,
     fp8_format,
+    lora_rank,
     target_type,
     new_repo_id,
     hf_token,
@@ -154,8 +174,10 @@ def process_and_upload_fp8(
             source_type, repo_url, safetensors_filename, hf_token, 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)
+        progress(0.25, desc="Converting to FP8 with LoRA extraction...")
+        success, msg = convert_safetensors_to_fp8_with_lora(
+            safetensors_path, output_dir, fp8_format, lora_rank, progress
+        )
         if not success:
             return None, f"❌ Conversion failed: {msg}", ""
 
@@ -165,42 +187,47 @@ def process_and_upload_fp8(
         )
 
         base_name = os.path.splitext(safetensors_filename)[0]
+        lora_filename = f"{base_name}-lora-r{lora_rank}.safetensors"
         readme = f"""---
 library_name: diffusers
 tags:
 - fp8
 - safetensors
-- delta-compensation
+- lora
+- low-rank
 - diffusion
 - converted-by-gradio
 ---
 
-# FP8 Model with Delta Compensation
+# FP8 Model with Low-Rank LoRA
 
 - **Source**: `{repo_url}`
 - **File**: `{safetensors_filename}`
 - **FP8 Format**: `{fp8_format.upper()}`
-- **Delta File**: `{base_name}-fp8-delta.safetensors`
+- **LoRA Rank**: {lora_rank}
+- **LoRA File**: `{lora_filename}`
 
 ## Usage (Inference)
 
-To restore near-original precision:
-
 ```python
-import torch
 from safetensors.torch import load_file
+import torch
 
+# Load FP8 model
 fp8_state = load_file("{base_name}-fp8-{fp8_format}.safetensors")
-delta_state = load_file("{base_name}-fp8-delta.safetensors")
+lora_state = load_file("{lora_filename}")
 
-restored_state = {{}}
+# Reconstruct approximate original weights
+reconstructed = {{}}
 for key in fp8_state:
-    if f"delta.{{key}}" in delta_state:
+    if f"lora_A.{{key}}" in lora_state and f"lora_B.{{key}}" in lora_state:
+        A = lora_state[f"lora_A.{{key}}"].to(torch.float32)
+        B = lora_state[f"lora_B.{{key}}"].to(torch.float32)
+        lora_weight = B @ A  # (rank, out) @ (in, rank) -> (out, in)
         fp8_weight = fp8_state[key].to(torch.float32)
-        delta = delta_state[f"delta.{{key}}"]
-        restored_state[key] = fp8_weight + delta
+        reconstructed[key] = fp8_weight + lora_weight
     else:
-        restored_state[key] = fp8_state[key].to(torch.float32)
+        reconstructed[key] = fp8_state[key].to(torch.float32)
 ```
 
 > Requires PyTorch ≥ 2.1 for FP8 support.
@@ -221,9 +248,9 @@ for key in fp8_state:
         result_html = f"""
 ✅ Success!  
 Model uploaded to: <a href="{repo_url_final}" target="_blank">{new_repo_id}</a>  
-Includes: FP8 model + delta compensation file.
+Includes: FP8 model + rank-{lora_rank} LoRA.
 """
-        return gr.HTML(result_html), "✅ FP8 + delta upload successful!", ""
+        return gr.HTML(result_html), "✅ FP8 + LoRA upload successful!", ""
 
     except Exception as e:
         return None, f"❌ Error: {str(e)}", ""
@@ -232,9 +259,9 @@ Includes: FP8 model + delta compensation file.
             shutil.rmtree(temp_dir, ignore_errors=True)
         shutil.rmtree(output_dir, ignore_errors=True)
 
-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.Blocks(title="FP8 + LoRA Extractor (HF ↔ ModelScope)") as demo:
+    gr.Markdown("# 🔄 FP8 Pruner with Low-Rank LoRA Extraction")
+    gr.Markdown("Convert `.safetensors` → **FP8** + **compact LoRA** for precision recovery. Supports Hugging Face ↔ ModelScope.")
 
     with gr.Row():
         with gr.Column():
@@ -242,6 +269,7 @@ with gr.Blocks(title="FP8 + Delta Converter (HF ↔ ModelScope)") as demo:
             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")
+            lora_rank = gr.Slider(minimum=8, maximum=256, step=8, value=64, label="LoRA Rank")
             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():
@@ -260,6 +288,7 @@ with gr.Blocks(title="FP8 + Delta Converter (HF ↔ ModelScope)") as demo:
             repo_url,
             safetensors_filename,
             fp8_format,
+            lora_rank,
             target_type,
             new_repo_id,
             hf_token,
@@ -272,9 +301,9 @@ with gr.Blocks(title="FP8 + Delta Converter (HF ↔ ModelScope)") as demo:
 
     gr.Examples(
         examples=[
-            ["huggingface", "https://huggingface.co/Yabo/FramePainter/tree/main", "unet_diffusion_pytorch_model.safetensors", "e5m2", "modelscope"]
+            ["huggingface", "https://huggingface.co/Yabo/FramePainter/tree/main", "unet_diffusion_pytorch_model.safetensors", "e5m2", 64, "modelscope"]
         ],
-        inputs=[source_type, repo_url, safetensors_filename, fp8_format, target_type]
+        inputs=[source_type, repo_url, safetensors_filename, fp8_format, lora_rank, target_type]
     )
 
 demo.launch()