Spaces:
Running on Zero
Running on Zero
| import os | |
| import gradio as gr | |
| import spaces | |
| from huggingface_hub import hf_hub_download, HfApi | |
| from convert_to_quant import quantize | |
| def do_quantize(quant_args): | |
| quantize(**quant_args) | |
| def run_quantization( | |
| source_repo, | |
| source_file, | |
| target_repo, | |
| target_filename_base, | |
| quant_format, | |
| layer_filter, | |
| exclude_layers_regex, | |
| full_precision_matrix_mult, | |
| generic_text, | |
| hf_token | |
| ): | |
| if not all([source_repo, source_file, target_repo, target_filename_base]): | |
| yield "Please fill in all repository and filename fields.", gr.update(visible=False), gr.update(visible=False) | |
| return | |
| try: | |
| # Download | |
| yield f"Downloading {source_file} from {source_repo}...", gr.update(visible=False), gr.update(visible=False) | |
| # We use token for download if provided, otherwise anonymous | |
| effective_token = hf_token if hf_token else os.getenv("HF_TOKEN") | |
| local_input_path = hf_hub_download(repo_id=source_repo, filename=source_file, token=effective_token if effective_token else None) | |
| # Setup quant arguments based on UI | |
| quant_args = { | |
| "input": local_input_path, | |
| "comfy_quant": True, | |
| "save_quant_metadata": True, | |
| "low_memory": False, | |
| "calib_cpu": True, | |
| "simple": True, | |
| "calib_samples": 8192 | |
| } | |
| suffix = "" | |
| if quant_format == "int8 rowwise": | |
| quant_args["int8"] = True | |
| quant_args["scaling_mode"] = "row" | |
| suffix = "-int8mixedrow-simple" | |
| elif quant_format == "int8-convrot": | |
| quant_args["int8"] = True | |
| quant_args["scaling_mode"] = "row" | |
| quant_args["convrot"] = True | |
| quant_args["convrot_group_size"] = 256 | |
| suffix = "-int8-convrot-simple" | |
| elif quant_format == "mxfp8": | |
| quant_args["mxfp8"] = True | |
| suffix = "-mxfp8mixed-simple" | |
| else: # fp8 (default) | |
| quant_args["scaling_mode"] = "tensor" | |
| suffix = "-fp8mixed-simple" | |
| output_filename = f"{target_filename_base}{suffix}.safetensors" | |
| output_path = f"./{output_filename}" | |
| quant_args["output"] = output_path | |
| # Layer filters | |
| if layer_filter == "Anima": quant_args["anima"] = True | |
| elif layer_filter == "Microsoft Lens": quant_args["lens"] = True | |
| elif layer_filter == "Flux2": quant_args["flux2"] = True | |
| elif layer_filter == "Chroma": quant_args["distillation_large"] = True | |
| elif layer_filter == "Radiance": quant_args["nerf_large"] = True; quant_args["radiance"] = True | |
| elif layer_filter == "WAN": quant_args["wan"] = True | |
| elif layer_filter == "LTX-2.x": quant_args["ltxv2"] = True | |
| elif layer_filter == "Qwen Image": quant_args["qwen"] = True | |
| elif layer_filter == "Z-Image": quant_args["zimage"] = True; quant_args["zimage_refiner"] = True | |
| elif layer_filter == "Krea2": quant_args["krea2"] = True | |
| elif layer_filter == "Boogu": quant_args["boogu"] = True | |
| elif layer_filter == "Ideogram4": quant_args["ideogram4"] = True | |
| if full_precision_matrix_mult: | |
| quant_args["full_precision_matrix_mult"] = True | |
| if generic_text: | |
| quant_args["generic_text"] = True | |
| if exclude_layers_regex: | |
| quant_args["exclude_layers"] = exclude_layers_regex | |
| yield f"Quantizing to {output_filename}...\nThis may take a few minutes.", gr.update(visible=False), gr.update(visible=False) | |
| do_quantize(quant_args) | |
| if effective_token: | |
| yield f"Uploading {output_filename} to {target_repo}...", gr.update(visible=False), gr.update(visible=False) | |
| # Upload | |
| try: | |
| api = HfApi(token=effective_token) | |
| commit_info = api.upload_file( | |
| path_or_fileobj=output_path, | |
| path_in_repo=output_filename, | |
| repo_id=target_repo, | |
| commit_message=f"Add {output_filename} quantized model", | |
| create_pr=True | |
| ) | |
| pr_url = commit_info.pr_url if hasattr(commit_info, 'pr_url') else f"https://huggingface.co/{target_repo}" | |
| yield ( | |
| f"Complete! Uploaded to {target_repo} and ready for direct download.", | |
| gr.update(value=f"<a href='{pr_url}' target='_blank' style='color: #3b82f6; text-decoration: underline; font-weight: bold;'>Click here to view the Pull Request</a>", visible=True), | |
| gr.update(value=output_path, visible=True) | |
| ) | |
| except Exception as upload_error: | |
| yield ( | |
| f"Complete! Upload failed: {str(upload_error)}. Ready for direct download below.", | |
| gr.update(visible=False), | |
| gr.update(value=output_path, visible=True) | |
| ) | |
| else: | |
| yield ( | |
| f"Complete! Ready for download below.", | |
| gr.update(visible=False), | |
| gr.update(value=output_path, visible=True) | |
| ) | |
| except Exception as e: | |
| yield f"Error: {str(e)}", gr.update(visible=False), gr.update(visible=False) | |
| # Build UI | |
| with gr.Blocks() as demo: | |
| with gr.Row(elem_id="topbar"): | |
| gr.Markdown("## 🤗 Model Quantizer", elem_classes=["brand"]) | |
| with gr.Row(elem_id="main-row", equal_height=True): | |
| with gr.Column(scale=4, min_width=280, elem_id="input-panel"): | |
| gr.Markdown("### Authentication (Optional)") | |
| hf_token = gr.Textbox(label="HF Token (WRITE)", type="password", placeholder="Paste your WRITE token for PR upload") | |
| no_token_checkbox = gr.Checkbox(label="I don't have a token", value=False) | |
| gr.Markdown("*If no token is provided, the environment variable option and direct download safeguard will be used.*", elem_classes=["text-sm"]) | |
| gr.Markdown("### Input Model") | |
| source_repo = gr.Textbox(label="Source HF Repo (e.g. author/model)") | |
| source_file = gr.Textbox(label="Source Filename (e.g. model.safetensors)") | |
| gr.Markdown("### Output Target") | |
| target_repo = gr.Textbox(label="Target HF Repo (e.g. author/model)") | |
| target_file_base = gr.Textbox(label="Target Filename Base (e.g. model-quant)") | |
| gr.Markdown("### Quantization Options") | |
| quant_format = gr.Radio(["fp8", "int8 rowwise", "int8-convrot", "mxfp8"], value="fp8", label="Format") | |
| layer_filter = gr.Dropdown( | |
| ["None", "Anima", "Microsoft Lens", "Flux2", "Chroma", "Radiance", "WAN", "LTX-2.x", "Qwen Image", "Z-Image", "Krea2", "Boogu", "Ideogram4"], | |
| value="None", label="Model Layer Filter" | |
| ) | |
| full_precision = gr.Checkbox(label="Full precision matrix multiplication", value=False) | |
| generic_text = gr.Checkbox(label="Generic text model quantization", value=False) | |
| exclude_layers = gr.Textbox(label="Exclude Layers Regex (Optional)", placeholder="(substring_1|substring_2)") | |
| run_btn = gr.Button("Quantize Model", variant="primary", size="lg") | |
| gr.Markdown("ℹ️ *For more advanced quantization modes, install and use [convert-to-quant](https://pypi.org/project/convert-to-quant/) locally.*", elem_classes=["text-sm", "mt-4"]) | |
| with gr.Column(scale=8, elem_id="output-panel"): | |
| status_text = gr.Textbox(label="Status Log", lines=10, interactive=False) | |
| output_link = gr.HTML(visible=False) | |
| output_file = gr.DownloadButton(label="Download Quantized Model", visible=False) | |
| def toggle_token_input(no_token): | |
| if no_token: | |
| return gr.update(interactive=False, value="") | |
| else: | |
| return gr.update(interactive=True) | |
| no_token_checkbox.change( | |
| fn=toggle_token_input, | |
| inputs=[no_token_checkbox], | |
| outputs=[hf_token] | |
| ) | |
| run_btn.click( | |
| fn=run_quantization, | |
| inputs=[ | |
| source_repo, source_file, target_repo, target_file_base, | |
| quant_format, layer_filter, exclude_layers, full_precision, | |
| generic_text, | |
| hf_token | |
| ], | |
| outputs=[status_text, output_link, output_file] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(css_paths=["assets/responsive.css"], theme=gr.themes.Default(primary_hue="blue", neutral_hue="zinc")) | |