Update app.py

app.py

@@ -8,13 +8,16 @@ from safetensors.torch import load_file, save_file
 
 TEMP_DIR = "temp_processing_dir"
 
-def convert_and_upload(token, source_repo, target_repo, precision):
+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."
         return
     if not target_repo.strip() or "your-username" in target_repo:
         yield "❌ Error: Please specify a valid Target Repository (e.g., your-username/repo-name)."
         return
+    if not target_components:
+        yield "❌ Error: Please select at least one component to quantize."
+        return
 
     # Map precision string to PyTorch dtype
     if precision == "FP8":
@@ -56,8 +59,11 @@ def convert_and_upload(token, source_repo, target_repo, precision):
                 local_dir_use_symlinks=False 
             )
 
-            # Check if it's a target safetensor file
-            if file.endswith(".safetensors") and ("text_encoder/" in file or "transformer/" in file):
+            # 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)
@@ -120,7 +126,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 Z-Image Quantizer & Uploader")
     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 **text_encoder** and **transformer** files, and uploads everything. "
+        "**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."
     )
 
@@ -140,6 +146,15 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 value="Tongyi-MAI/Z-Image-Turbo", 
                 label="Source Repository"
             )
+            
+            # 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."
+            )
+            
             precision = gr.Dropdown(
                 choices=["FP8", "FP16", "BF16"], 
                 value="FP8", 
@@ -171,7 +186,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     start_btn.click(
         fn=convert_and_upload, 
-        inputs=[hf_token, source_repo, target_repo, precision], 
+        inputs=[hf_token, source_repo, target_repo, precision, target_components], 
         outputs=[output_log]
     )