import os import gradio as gr import spaces from huggingface_hub import hf_hub_download, HfApi from convert_to_quant import quantize @spaces.GPU(duration=300) 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"Click here to view the Pull Request", 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"))