Update app.py

app.py

@@ -2,12 +2,11 @@ import os
 import gc
 import torch
 import shutil
+import uuid
 import gradio as gr
 from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file, save_file
 
-TEMP_DIR = "temp_processing_dir"
-
 def convert_and_upload(token, source_repo, target_repo, precision, target_components):
     if not token:
         yield "❌ Error: Please provide a valid Hugging Face Write Token."
@@ -19,15 +18,13 @@ def convert_and_upload(token, source_repo, target_repo, precision, target_compon
         yield "❌ Error: Please select at least one component to quantize."
         return
 
-    # Map precision string to PyTorch dtype
-    if precision == "FP8":
-        target_dtype = torch.float8_e4m3fn
-    elif precision == "FP16":
-        target_dtype = torch.float16
-    elif precision == "BF16":
-        target_dtype = torch.bfloat16
-    else:
-        target_dtype = None
+    target_dtype = None
+    is_int8 = False
+    
+    if precision == "FP8": target_dtype = torch.float8_e4m3fn
+    elif precision == "FP16": target_dtype = torch.float16
+    elif precision == "BF16": target_dtype = torch.bfloat16
+    elif precision == "INT8": is_int8 = True
 
     api = HfApi(token=token)
     yield f"🔄 Connecting to Hugging Face and verifying target repo: {target_repo}..."
@@ -45,41 +42,81 @@ def convert_and_upload(token, source_repo, target_repo, precision, target_compon
         yield f"❌ Error fetching files: {str(e)}"
         return
 
-    os.makedirs(TEMP_DIR, exist_ok=True)
+    cache_dir = f"./hf_cache_{uuid.uuid4().hex[:8]}"
+    success_count = 0
+    error_count = 0
+
+    # Z-IMAGE SPECIFIC EXCLUSIONS
+    # Protects the DiT's embedders/final layers and the Text Encoder's sensitive norms from INT8 destruction
+    exclude_prefixes = [
+        "t_embedder", "cap_embedder", "all_x_embedder", "all_final_layer", "rope_embedder", 
+        "embed_tokens", "norm", "ln_", "shared"
+    ]
 
     for file in files:
+        is_root_safetensor = "/" not in file and file.endswith(".safetensors")
+        
+        if is_root_safetensor:
+            yield f"🗑️ Auto-skipping root model: {file}..."
+            try:
+                api.delete_file(path_in_repo=file, repo_id=target_repo, token=token, commit_message=f"Auto-deleted {file}")
+            except Exception:
+                pass 
+            continue
+
         yield f"⏳ Processing {file}..."
 
         try:
-            # Download file locally, bypassing symlink cache to save disk space
+            os.makedirs(cache_dir, exist_ok=True)
+
             local_path = hf_hub_download(
                 repo_id=source_repo,
                 filename=file,
-                local_dir=TEMP_DIR,
-                local_dir_use_symlinks=False 
+                cache_dir=cache_dir,
+                token=token 
             )
 
-            # Check if this file belongs to one of the selected target components
             in_target_component = any(f"{comp}/" in file for comp in target_components)
 
-            # Intercept and quantize only if it's a safetensors file in a selected folder
             if file.endswith(".safetensors") and in_target_component:
                 yield f"🧠 Quantizing {file} to {precision}..."
                 
                 tensors = load_file(local_path)
+                new_tensors = {}
+
+                for k, v in tensors.items():
+                    # --- BRANCH 1: INT8 Symmetric Quantization ---
+                    if is_int8:
+                        is_2d_weight = "weight" in k and len(v.shape) == 2
+                        is_excluded = any(ex in k for ex in exclude_prefixes)
+
+                        if is_2d_weight and not is_excluded:
+                            # Upcast to BF16 for math
+                            if v.dtype == torch.float8_e4m3fn:
+                                v = v.to(torch.bfloat16)
+
+                            scale = v.abs().max(dim=1, keepdim=True)[0] / 127.0
+                            scale = scale.clamp(min=1e-8)
+                            weight_int8 = torch.round(v / scale).clamp(-127, 127).to(torch.int8)
+                            
+                            base_name = k.rsplit(".", 1)[0]
+                            new_tensors[f"{base_name}.weight_int8"] = weight_int8
+                            new_tensors[f"{base_name}.weight_scale"] = scale.to(torch.bfloat16)
+                        else:
+                            new_tensors[k] = v.to(torch.bfloat16) if v.is_floating_point() else v
+
+                    # --- BRANCH 2: Standard Floating Point Casting ---
+                    else:
+                        if v.is_floating_point():
+                            new_tensors[k] = v.to(target_dtype)
+                        else:
+                            new_tensors[k] = v
+
+                converted_path = "converted.safetensors"
+                save_file(new_tensors, converted_path)
 
-                # Cast floating point tensors to the selected precision
-                if target_dtype:
-                    keys = list(tensors.keys())
-                    for k in keys:
-                        if tensors[k].is_floating_point():
-                            tensors[k] = tensors[k].to(target_dtype)
-
-                converted_path = os.path.join(TEMP_DIR, "converted.safetensors")
-                save_file(tensors, converted_path)
-
-                # Wipe tensors from RAM to prevent OOM
                 del tensors
+                del new_tensors
                 gc.collect()
 
                 yield f"☁️ Uploading {precision} version of {file}..."
@@ -101,88 +138,76 @@ def convert_and_upload(token, source_repo, target_repo, precision, target_compon
                     commit_message=f"Copy {file} from original repo"
                 )
 
-            # Cleanup original downloaded file
-            if os.path.exists(local_path):
-                os.remove(local_path)
-            
+            success_count += 1
+
+            if os.path.exists(cache_dir):
+                shutil.rmtree(cache_dir)
             gc.collect()
 
         except Exception as e:
-            yield f"⚠️ Error processing {file}: {str(e)}\nSkipping to next file..."
+            error_count += 1
+            yield f"⚠️ Error processing {file}: {str(e)}\nSkipping..."
+
+    if os.path.exists(cache_dir):
+        shutil.rmtree(cache_dir)
 
-    if os.path.exists(TEMP_DIR):
-        shutil.rmtree(TEMP_DIR)
+    yield f"✅ Finished! Successfully processed {success_count} files. Errors encountered: {error_count}."
 
-    yield f"✅ All files processed and successfully uploaded to {target_repo}!"
 
-# Dynamic UI Update for Target Repo Name
 def update_target_repo(username, source, precision):
     user_prefix = username.strip() if username.strip() else "your-username"
-    model_name = "Z-Image-Turbo" if "Turbo" in source else "Z-Image-Base"
+    model_name = source.split("/")[-1] if "/" in source else source
     return f"{user_prefix}/{model_name}-{precision}"
 
-# Build the Gradio UI
+def update_warnings(precision):
+    if precision == "INT8":
+        return gr.update(value="⚠️ **INT8 Warning:** Modifies layer keys (`weight_int8`, `weight_scale`). Requires the custom `NativeInt8Linear` XPU inference code to run.", visible=True)
+    else:
+        return gr.update(visible=False)
+
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Z-Image Quantizer & Uploader")
+    gr.Markdown("# 🚀 Universal Z-Image Quantizer")
     gr.Markdown(
-        "Convert the **Z-Image** or **Z-Image-Turbo** models to lower precisions (FP8, FP16, BF16) and push them directly to your own Hugging Face account.\n\n"
-        "**How it works:** This tool sequentially downloads, quantizes the selected files, and uploads everything. "
-        "It is designed to run safely on free Spaces (16GB RAM) by processing files one at a time."
+        "Convert sharded Z-Image models directly on Hugging Face to floating-point precisions (FP8/FP16/BF16) or dynamically trigger symmetric integer quantization (INT8)."
     )
 
     with gr.Row():
         with gr.Column(scale=2):
-            hf_token = gr.Textbox(
-                label="Hugging Face Token (Write Access Required)", 
-                type="password",
-                placeholder="hf_..."
-            )
-            hf_username = gr.Textbox(
-                label="Your Hugging Face Username", 
-                placeholder="e.g., rootlocalghost"
-            )
+            hf_token = gr.Textbox(label="Hugging Face Token (Write Access)", type="password", placeholder="hf_...")
+            hf_username = gr.Textbox(label="Hugging Face Username", placeholder="e.g., rootlocalghost")
+            
             source_repo = gr.Dropdown(
-                choices=["Tongyi-MAI/Z-Image", "Tongyi-MAI/Z-Image-Turbo"], 
-                value="Tongyi-MAI/Z-Image-Turbo", 
-                label="Source Repository"
+                choices=["your-username/Z-Image-Turbo", "your-username/Z-Image-Base"], 
+                value="your-username/Z-Image-Turbo", 
+                label="Source Repository", 
+                allow_custom_value=True
             )
             
-            # Added checkbox group for granular component control
             target_components = gr.CheckboxGroup(
-                choices=["text_encoder", "transformer"],
-                value=["text_encoder", "transformer"],
-                label="Components to Quantize",
-                info="Select which parts of the model to convert. Unselected parts will be copied as-is."
+                choices=["text_encoder", "transformer", "vae"],
+                value=["transformer"],
+                label="Components to Quantize"
             )
             
             precision = gr.Dropdown(
-                choices=["FP8", "FP16", "BF16"], 
-                value="FP8", 
-                label="Quantization Precision"
-            )
-            target_repo = gr.Textbox(
-                label="Target Repository (Auto-generated)", 
-                value="your-username/Z-Image-Turbo-FP8",
-                interactive=True
+                choices=["FP8", "FP16", "BF16", "INT8"], 
+                value="INT8", 
+                label="Target Precision"
             )
+            
+            int8_warning = gr.Markdown(visible=True, value="⚠️ **INT8 Warning:** Modifies layer keys (`weight_int8`, `weight_scale`). Requires the custom `NativeInt8Linear` XPU inference code to run.")
+            
+            target_repo = gr.Textbox(label="Target Repository (Auto-generated)", value="your-username/Z-Image-Turbo-INT8", interactive=True)
             start_btn = gr.Button("Start Quantization & Upload", variant="primary")
         
         with gr.Column(scale=3):
-            output_log = gr.Textbox(
-                label="Operation Logs", 
-                lines=17, 
-                interactive=False,
-                max_lines=20
-            )
+            output_log = gr.Textbox(label="Operation Logs", lines=20, interactive=False, max_lines=25)
 
-    # Automatically update the target repo name when inputs change
     inputs_to_watch = [hf_username, source_repo, precision]
     for inp in inputs_to_watch:
-        inp.change(
-            fn=update_target_repo, 
-            inputs=inputs_to_watch, 
-            outputs=[target_repo]
-        )
+        inp.change(fn=update_target_repo, inputs=inputs_to_watch, outputs=[target_repo])
+        
+    precision.change(fn=update_warnings, inputs=[precision], outputs=[int8_warning])
 
     start_btn.click(
         fn=convert_and_upload,