core/generation_logic.py

from typing import Any, Dict
import gradio as gr

from core.pipelines.controlnet_preprocessor import ControlNetPreprocessorPipeline
from core.pipelines.sd_image_pipeline import SdImagePipeline

controlnet_preprocessor_pipeline = ControlNetPreprocessorPipeline()
sd_image_pipeline = SdImagePipeline()


def build_reverse_map():
    from nodes import NODE_DISPLAY_NAME_MAPPINGS
    import core.pipelines.controlnet_preprocessor as cn_module
    
    if cn_module.REVERSE_DISPLAY_NAME_MAP is None:
        cn_module.REVERSE_DISPLAY_NAME_MAP = {v: k for k, v in NODE_DISPLAY_NAME_MAPPINGS.items()}
        if "Semantic Segmentor (legacy, alias for UniFormer)" not in cn_module.REVERSE_DISPLAY_NAME_MAP:
             cn_module.REVERSE_DISPLAY_NAME_MAP["Semantic Segmentor (legacy, alias for UniFormer)"] = "SemSegPreprocessor"


def run_cn_preprocessor_entry(*args, **kwargs):
    return controlnet_preprocessor_pipeline.run(*args, **kwargs)

def generate_image_wrapper(ui_inputs: dict, progress=gr.Progress(track_tqdm=True)):
    return sd_image_pipeline.run(ui_inputs=ui_inputs, progress=progress)