app.py

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"<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"))