Gin3601-LLS-node/qwen/nodes.py

"""
High-level Qwen nodes.
"""
from __future__ import annotations

import importlib

from ..sampling.nodes import _get_samplers, _get_schedulers
from . import discovery, runtime

try:
    comfy_core_nodes = importlib.import_module("nodes")
except Exception:
    comfy_core_nodes = None

_MAX_RESOLUTION = int(getattr(comfy_core_nodes, "MAX_RESOLUTION", 8192))


class LLSQwenTextToImage:
    CATEGORY = "LLS/Qwen"
    FUNCTION = "generate"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    DESCRIPTION = (
        "Compressed Qwen text-to-image node. Internally loads Qwen resources when needed, "
        "supports advanced sampling controls, optional external MODEL chaining, "
        "and optional turbo/lightning LoRAs."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (discovery.get_qwen_text_model_choices(),),
                "prompt": ("STRING", {"default": "", "multiline": True}),
                "width": ("INT", {"default": 1024, "min": 16, "max": _MAX_RESOLUTION, "step": 16}),
                "height": ("INT", {"default": 1024, "min": 16, "max": _MAX_RESOLUTION, "step": 16}),
                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
                "negative_prompt": ("STRING", {"default": "", "multiline": True, "advanced": True}),
                "cfg": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step": 0.1, "advanced": True}),
                "sampler_name": (_get_samplers(), {"default": "euler", "advanced": True}),
                "scheduler": (_get_schedulers(), {"default": "simple", "advanced": True}),
                "shift": ("FLOAT", {"default": 3.1, "min": 0.0, "max": 100.0, "step": 0.01, "advanced": True}),
                "enable_turbo_mode": ("BOOLEAN", {"default": False, "advanced": True}),
                "turbo_lora_name": (
                    discovery.get_qwen_text_turbo_lora_choices(),
                    {"default": discovery.AUTO_TURBO_LORA_CHOICE, "advanced": True},
                ),
                "turbo_strength": (
                    "FLOAT",
                    {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "advanced": True},
                ),
            },
            "optional": {
                "model": ("MODEL",),
            },
        }

    def generate(
        self,
        model_name: str,
        prompt: str,
        width: int,
        height: int,
        steps: int,
        seed: int,
        batch_size: int,
        negative_prompt: str,
        cfg: float,
        sampler_name: str,
        scheduler: str,
        shift: float,
        enable_turbo_mode: bool,
        turbo_lora_name: str,
        turbo_strength: float,
        model=None,
    ):
        return (
            runtime.run_qwen_text_to_image(
                model_name=model_name,
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                steps=steps,
                seed=seed,
                batch_size=batch_size,
                cfg=cfg,
                sampler_name=sampler_name,
                scheduler=scheduler,
                shift=shift,
                enable_turbo_mode=enable_turbo_mode,
                turbo_lora_name=turbo_lora_name,
                turbo_strength=turbo_strength,
                model=model,
            ),
        )


class LLSQwenImageEdit:
    CATEGORY = "LLS/Qwen"
    FUNCTION = "generate"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    DESCRIPTION = (
        "Compressed Qwen image-edit node. Internally loads Qwen companion resources, "
        "supports official multi-image edit conditioning, optional external MODEL chaining, "
        "and optional turbo/lightning LoRAs."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (discovery.get_qwen_edit_model_choices(),),
                "image": ("IMAGE",),
                "prompt": ("STRING", {"default": "", "multiline": True}),
                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
                "negative_prompt": ("STRING", {"default": "", "multiline": True, "advanced": True}),
                "cfg": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step": 0.1, "advanced": True}),
                "sampler_name": (_get_samplers(), {"default": "euler", "advanced": True}),
                "scheduler": (_get_schedulers(), {"default": "simple", "advanced": True}),
                "shift": ("FLOAT", {"default": 3.1, "min": 0.0, "max": 100.0, "step": 0.01, "advanced": True}),
                "cfg_norm_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "advanced": True}),
                "reference_latents_method": (
                    discovery.REFERENCE_LATENTS_METHOD_CHOICES,
                    {"default": "index_timestep_zero", "advanced": True},
                ),
                "enable_turbo_mode": ("BOOLEAN", {"default": False, "advanced": True}),
                "turbo_lora_name": (
                    discovery.get_qwen_edit_turbo_lora_choices(),
                    {"default": discovery.AUTO_TURBO_LORA_CHOICE, "advanced": True},
                ),
                "turbo_strength": (
                    "FLOAT",
                    {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "advanced": True},
                ),
            },
            "optional": {
                "image2": ("IMAGE",),
                "image3": ("IMAGE",),
                "model": ("MODEL",),
            },
        }

    def generate(
        self,
        model_name: str,
        image,
        prompt: str,
        steps: int,
        seed: int,
        image2=None,
        image3=None,
        negative_prompt: str = "",
        cfg: float = 4.0,
        sampler_name: str = "euler",
        scheduler: str = "simple",
        shift: float = 3.1,
        cfg_norm_strength: float = 1.0,
        reference_latents_method: str = "index_timestep_zero",
        enable_turbo_mode: bool = False,
        turbo_lora_name: str = discovery.AUTO_TURBO_LORA_CHOICE,
        turbo_strength: float = 1.0,
        model=None,
    ):
        return (
            runtime.run_qwen_image_edit(
                model_name=model_name,
                image=image,
                image2=image2,
                image3=image3,
                prompt=prompt,
                negative_prompt=negative_prompt,
                steps=steps,
                seed=seed,
                cfg=cfg,
                sampler_name=sampler_name,
                scheduler=scheduler,
                shift=shift,
                cfg_norm_strength=cfg_norm_strength,
                reference_latents_method=reference_latents_method,
                enable_turbo_mode=enable_turbo_mode,
                turbo_lora_name=turbo_lora_name,
                turbo_strength=turbo_strength,
                model=model,
            ),
        )


NODE_CLASS_MAPPINGS: dict[str, type] = {
    "LLSQwenTextToImage": LLSQwenTextToImage,
    "LLSQwenImageEdit": LLSQwenImageEdit,
}

NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
    "LLSQwenTextToImage": "LLS Qwen Text To Image",
    "LLSQwenImageEdit": "LLS Qwen Image Edit",
}