ui/shared/ui_components.py

import gradio as gr
from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
from core.settings import (
    MAX_LORAS, MAX_CONTROLNETS, MAX_H3_GUIDES, LORA_SOURCE_CHOICES, RESOLUTION_MAP, MODEL_DEFAULTS_CONFIG
)
from utils.app_utils import get_h3_cn_defaults

default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
DEFAULT_STEPS = default_vals.get('steps', 20)
DEFAULT_CFG = default_vals.get('cfg', 5.0)
DEFAULT_SAMPLER = default_vals.get('sampler_name', 'euler')
DEFAULT_SCHEDULER = default_vals.get('scheduler', 'simple')
DEFAULT_POS_PROMPT = default_vals.get('positive_prompt', '')
DEFAULT_NEG_PROMPT = default_vals.get('negative_prompt', '')


def create_base_parameter_ui(prefix, defaults=None):
    if defaults is None:
        defaults = {}
    
    components = {}
    
    with gr.Row():
        components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
            label="Aspect Ratio",
            choices=list(RESOLUTION_MAP.get('sdxl', {}).keys()),
            value="1:1 (Square)",
            interactive=True,
            allow_custom_value=True
        )
    with gr.Row():
        components[f'width_{prefix}'] = gr.Number(label="Width", value=defaults.get('w', 1024), interactive=True)
        components[f'height_{prefix}'] = gr.Number(label="Height", value=defaults.get('h', 1024), interactive=True)
    with gr.Row():
        components[f'sampler_{prefix}'] = gr.Dropdown(
            label="Sampler", 
            choices=SAMPLER_CHOICES, 
            value=DEFAULT_SAMPLER if DEFAULT_SAMPLER in SAMPLER_CHOICES else (SAMPLER_CHOICES[0] if SAMPLER_CHOICES else 'euler')
        )
        components[f'scheduler_{prefix}'] = gr.Dropdown(
            label="Scheduler", 
            choices=SCHEDULER_CHOICES, 
            value=DEFAULT_SCHEDULER if DEFAULT_SCHEDULER in SCHEDULER_CHOICES else (SCHEDULER_CHOICES[0] if SCHEDULER_CHOICES else 'simple')
        )
    with gr.Row():
        components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=DEFAULT_STEPS)
        components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=DEFAULT_CFG)
    with gr.Row():
        components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
        components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
    with gr.Row():
        components[f'guidance_{prefix}'] = gr.Slider(label="Guidance (FLUX)", minimum=1.0, maximum=10.0, step=0.1, value=3.5, visible=False, interactive=True)
        components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60s, Max: 120s", info="Optional: Set how long to reserve the GPU.")
            
    return components


def create_lora_settings_ui(prefix: str):
    components = {}
    
    lora_rows, lora_sources, lora_ids, lora_scales, lora_uploads = [], [], [], [], []

    with gr.Accordion("LoRA Settings", open=False, visible=True) as lora_accordion:
        components[f'lora_accordion_{prefix}'] = lora_accordion
        gr.Markdown("💡 **Tip:** When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button. When downloading from Hugging Face, please use the format: `repo_id/filename.extension` or `repo_id/folder_path/filename.extension` (e.g., `Comfy-Org/MiniMax-H3/loras/minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors` or `Comfy-Org/MiniMax-H3/loras/minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors`).")
        components[f'lora_count_state_{prefix}'] = gr.State(1)

        for i in range(MAX_LORAS):
            with gr.Row(visible=i==0) as row:
                source = gr.Dropdown(label=f"LoRA Source {i+1}", choices=LORA_SOURCE_CHOICES, value=LORA_SOURCE_CHOICES[0], scale=1)
                lora_id = gr.Textbox(label="Civitai Version ID / HF file / Upload File", scale=2, type="text")
                scale = gr.Slider(label=f"Scale", minimum=0.0, maximum=2.0, step=0.05, value=1.0, scale=1)
                upload = gr.UploadButton(label="Upload", file_types=[".safetensors"], scale=1)
            
            lora_rows.append(row)
            lora_sources.append(source)
            lora_ids.append(lora_id)
            lora_scales.append(scale)
            lora_uploads.append(upload)
            
        with gr.Row():
            components[f'add_lora_button_{prefix}'] = gr.Button("Add LoRA", variant="secondary")
            components[f'delete_lora_button_{prefix}'] = gr.Button("Remove LoRA", variant="secondary", visible=False)

    components[f'lora_rows_{prefix}'] = lora_rows
    components[f'lora_sources_{prefix}'] = lora_sources
    components[f'lora_ids_{prefix}'] = lora_ids
    components[f'lora_scales_{prefix}'] = lora_scales
    components[f'lora_uploads_{prefix}'] = lora_uploads
    
    all_lora_components_flat = []
    for i in range(MAX_LORAS):
        all_lora_components_flat.extend([lora_sources[i], lora_ids[i], lora_scales[i], lora_uploads[i]])
    components[f'all_lora_components_flat_{prefix}'] = all_lora_components_flat
    
    return components


def create_h3_guide_ui(prefix: str, max_units=MAX_H3_GUIDES):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    with gr.Accordion("Keyframe Guide Settings", open=False, visible=True) as accordion:
        components[key('h3_guide_accordion')] = accordion
        gr.Markdown(
            "💡 **Tip:** You can anchor keyframe guides (image, video clip, or audio) "
            "at arbitrary positions along the continuous timeline (not limited to first/last frames). "
            "Specify the position via **Time (seconds)**."
        )

        guide_rows = []
        images = []
        videos = []
        audios = []
        times = []
        frames = []

        components.update({
            key('h3_guide_rows'): guide_rows,
            key('h3_guide_images'): images,
            key('h3_guide_videos'): videos,
            key('h3_guide_audios'): audios,
            key('h3_guide_times'): times,
            key('h3_guide_frames'): frames,
        })

        for i in range(max_units):
            with gr.Row(visible=(i < 1)) as row:
                with gr.Column(scale=1):
                    images.append(gr.Image(label=f"Guide Image {i+1}", type="pil", sources=["upload"], height=200))
                with gr.Column(scale=1):
                    videos.append(gr.Video(label=f"Guide Video (Opt) {i+1}", sources=["upload"], height=200))
                with gr.Column(scale=1):
                    audios.append(gr.Audio(label=f"Guide Audio (Opt) {i+1}", sources=["upload"], type="filepath"))
                    times.append(gr.Slider(label=f"Time (s) {i+1}", minimum=0.0, maximum=15.0, step=0.5, value=float(i * 1.5), interactive=True))
                guide_rows.append(row)

        with gr.Row():
            components[key('add_h3_guide_button')] = gr.Button("✚ Add Guide", variant="secondary")
            components[key('delete_h3_guide_button')] = gr.Button("➖ Delete Guide", variant="secondary", visible=False)
        components[key('h3_guide_count_state')] = gr.State(1)

        all_guide_components_flat = []
        for i in range(max_units):
            all_guide_components_flat.extend([
                images[i], videos[i], audios[i], times[i]
            ])
        components[key('all_h3_guide_components_flat')] = all_guide_components_flat

    return components


def create_h3_controlnet_ui(prefix: str, max_units=MAX_CONTROLNETS):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    default_types, default_type, default_series, default_serie, default_fp = get_h3_cn_defaults()

    with gr.Accordion("H3 ControlNet Settings", open=False, visible=True) as accordion:
        components[key('h3_controlnet_accordion')] = accordion
        gr.Markdown(
            "💡 **Tip:** This feature does not perform preprocessing on uploaded videos. "
            "Please upload the corresponding preprocessed control video (e.g., Canny edges, Depth, Pose) rather than raw/unprocessed video footage."
        )

        cn_rows, videos, series, types, strengths, start_percents, end_percents, filepaths = [], [], [], [], [], [], [], []
        components.update({
            key('h3_controlnet_rows'): cn_rows,
            key('h3_controlnet_videos'): videos,
            key('h3_controlnet_series'): series,
            key('h3_controlnet_types'): types,
            key('h3_controlnet_strengths'): strengths,
            key('h3_controlnet_start_percents'): start_percents,
            key('h3_controlnet_end_percents'): end_percents,
            key('h3_controlnet_filepaths'): filepaths,
        })

        for i in range(max_units):
            with gr.Row(visible=(i < 1)) as row:
                with gr.Column(scale=1):
                    videos.append(gr.Video(label=f"Control Video {i+1}", sources=["upload"], height=220))
                with gr.Column(scale=2):
                    types.append(gr.Dropdown(label="Type", choices=default_types, value=default_type, interactive=True, allow_custom_value=True))
                    series.append(gr.Dropdown(label="Series", choices=default_series, value=default_serie, interactive=True, allow_custom_value=True))
                    strengths.append(gr.Slider(label="Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
                    start_percents.append(gr.State(0.0))
                    end_percents.append(gr.State(1.0))
                    filepaths.append(gr.State(default_fp))
                cn_rows.append(row)

        with gr.Row():
            components[key('add_h3_controlnet_button')] = gr.Button("✚ Add H3 ControlNet", variant="secondary")
            components[key('delete_h3_controlnet_button')] = gr.Button("➖ Delete H3 ControlNet", variant="secondary", visible=False)
        components[key('h3_controlnet_count_state')] = gr.State(1)

        all_cn_components_flat = []
        for i in range(max_units):
            all_cn_components_flat.extend([
                videos[i], types[i], series[i], strengths[i], start_percents[i], end_percents[i], filepaths[i]
            ])
        components[key('all_h3_controlnet_components_flat')] = all_cn_components_flat

    return components