src/ui/loader_tab.py

"""Pipeline loader tab for SDXL Model Merger."""

import gradio as gr

from ..config import (
    DEFAULT_CHECKPOINT_URL,
    DEFAULT_VAE_URL,
    DEFAULT_LORA_URLS,
)
from ..pipeline import load_pipeline, cancel_download


def create_loader_tab():
    """Create the pipeline loading tab with all input controls."""

    with gr.Accordion("⚙️ 1. Load Pipeline", open=True, elem_classes=["feature-card"]):
        with gr.Row():
            with gr.Column(scale=2):
                # Checkpoint URL
                checkpoint_url = gr.Textbox(
                    label="Base Model (.safetensors) URL",
                    value=DEFAULT_CHECKPOINT_URL,
                    placeholder="e.g., https://civitai.com/api/download/models/...",
                    info="Download link for the base SDXL checkpoint"
                )

                # VAE URL (optional)
                vae_url = gr.Textbox(
                    label="VAE (.safetensors) URL",
                    value=DEFAULT_VAE_URL,
                    placeholder="Leave blank to use model's built-in VAE",
                    info="Optional custom VAE for improved quality"
                )

            with gr.Column(scale=1):
                # LoRA URLs
                lora_urls = gr.Textbox(
                    label="LoRA URLs (one per line)",
                    lines=5,
                    value=DEFAULT_LORA_URLS,
                    placeholder="https://civit.ai/...\nhttps://huggingface.co/...",
                    info="Multiple LoRAs can be loaded and fused together"
                )
                
                # LoRA strengths
                lora_strengths = gr.Textbox(
                    label="LoRA Strengths",
                    value="1.0",
                    placeholder="e.g., 0.8,1.0,0.5",
                    info="Comma-separated strength values for each LoRA"
                )
        
        # Action buttons
        with gr.Row():
            load_btn = gr.Button("🚀 Load Pipeline", variant="primary", size="lg")
            cancel_btn = gr.Button("🛑 Cancel Download", variant="stop", size="lg")
        
        # Status output
        load_status = gr.Textbox(
            label="Status",
            placeholder="Ready to load pipeline...",
            show_label=True,
        )
    
    return (
        checkpoint_url, vae_url, lora_urls, lora_strengths,
        load_btn, cancel_btn, load_status
    )


def setup_loader_events(
    checkpoint_url, vae_url, lora_urls, lora_strengths,
    load_btn, cancel_btn, load_status
):
    """Setup event handlers for the loader tab."""
    
    load_btn.click(
        fn=load_pipeline,
        inputs=[checkpoint_url, vae_url, lora_urls, lora_strengths],
        outputs=load_status,
    )
    
    cancel_btn.click(fn=cancel_download)