src/ui/exporter_tab.py

"""Model exporter tab for SDXL Model Merger."""

import gradio as gr

from ..exporter import export_merged_model


def create_exporter_tab():
    """Create the model export tab with all configuration options."""
    
    with gr.Accordion("📦 3. Export Merged Model", open=True, elem_classes=["feature-card"]):
        # Export settings
        with gr.Row():
            include_lora = gr.Checkbox(
                True,
                label="Include Fused LoRAs",
                info="Bake the loaded LoRAs into the exported model"
            )
            
            quantize_toggle = gr.Checkbox(
                False,
                label="Apply Quantization",
                info="Reduce model size with quantization"
            )
        
        # Quantization options
        with gr.Row(visible=True) as qtype_row:
            qtype_dropdown = gr.Dropdown(
                choices=["none", "int8", "int4", "float8"],
                value="int8",
                label="Quantization Method",
                info="Trade quality for smaller file size"
            )
        
        # Format options
        with gr.Row():
            format_dropdown = gr.Dropdown(
                choices=["safetensors", "bin"],
                value="safetensors",
                label="Export Format",
                info="safetensors is recommended for safety"
            )
        
        # Export button and output
        with gr.Row():
            export_btn = gr.Button("💾 Save Merged Checkpoint", variant="primary", size="lg")
        
        with gr.Row():
            download_link = gr.File(
                label="Download Merged File",
                show_label=True,
            )
            
            with gr.Column():
                export_status = gr.Textbox(
                    label="Export Status",
                    placeholder="Ready to export..."
                )
                
                # Info about quantization
                gr.HTML("""
                    <div style="margin-top: 16px; padding: 12px; background: #e0f2fe; border-radius: 8px;">
                        <strong>ℹ️ About Quantization:</strong>
                        <p style="font-size: 0.9em; margin: 8px 0;">
                            Reduces model size by lowering precision using torchao.
                            Int8 is typically lossless for inference while cutting size in half.
                            Int4 provides maximum compression with minimal quality loss.
                        </p>
                    </div>
                """)
    
    return (
        include_lora, quantize_toggle, qtype_dropdown, format_dropdown,
        export_btn, download_link, export_status, qtype_row
    )


def setup_exporter_events(
    include_lora, quantize_toggle, qtype_dropdown, format_dropdown,
    export_btn, download_link, export_status, qtype_row
):
    """Setup event handlers for the exporter tab."""
    
    # Toggle quantization row visibility
    quantize_toggle.change(
        fn=lambda checked: gr.update(visible=checked),
        inputs=[quantize_toggle],
        outputs=qtype_row,
    )
    
    # Clear download link after use
    def clear_download_link():
        return None
    
    export_btn.click(
        fn=lambda inc, q, qt, fmt: export_merged_model(
            include_lora=inc,
            quantize=q and (qt != "none"),
            qtype=qt if qt != "none" else None,
            save_format=fmt,
        ),
        inputs=[include_lora, quantize_toggle, qtype_dropdown, format_dropdown],
        outputs=[download_link, export_status],
    ).then(
        fn=clear_download_link,
        outputs=[download_link],
    )