AMXELA-Official-ComfyUI-SeFiImage/nodes.py

"""ComfyUI wrapper nodes for SeFi-Image (Semantic-First Diffusion).

Upstream model + inference package: https://github.com/jmliu206/SeFi-Image (MIT).
The `sefi` package is NOT pip-installable (no packaging metadata), so on first
use we fetch the repo source as a tarball (stdlib only, no git required) into
this pack directory and add it to sys.path. Heavy imports are lazy so that a
missing dependency or network hiccup never blocks ComfyUI startup.
"""

import io
import os
import sys
import shutil
import tarfile
import tempfile
import urllib.request

import numpy as np
import torch

PACK_DIR = os.path.dirname(os.path.abspath(__file__))
VENDOR_DIR = os.path.join(PACK_DIR, "SeFi-Image")
SEFI_TARBALL = "https://codeload.github.com/jmliu206/SeFi-Image/tar.gz/refs/heads/main"

# The HF checkpoint ids published under the SeFi-Image org.
CHECKPOINTS = [
    "SeFi-Image/SeFi-Image-2B-turbo",
    "SeFi-Image/SeFi-Image-1B-turbo",
    "SeFi-Image/SeFi-Image-5B-turbo",
    "SeFi-Image/SeFi-Image-1B-Base",
    "SeFi-Image/SeFi-Image-2B-Base",
    "SeFi-Image/SeFi-Image-5B-Base",
    "SeFi-Image/SeFi-Image-5B-RL",
]

_DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}


def _fetch_sefi_source():
    """Download + extract the SeFi-Image repo into VENDOR_DIR (no git needed)."""
    if os.path.isdir(os.path.join(VENDOR_DIR, "sefi")):
        return
    tmp = tempfile.mkdtemp(prefix="sefi_dl_")
    try:
        data = urllib.request.urlopen(SEFI_TARBALL, timeout=120).read()
        with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tf:
            tf.extractall(tmp)
        # Tarball extracts to a single top-level dir like "SeFi-Image-main".
        roots = [d for d in os.listdir(tmp) if os.path.isdir(os.path.join(tmp, d))]
        if not roots:
            raise RuntimeError("SeFi-Image tarball had no top-level directory")
        src = os.path.join(tmp, roots[0])
        if os.path.isdir(VENDOR_DIR):
            shutil.rmtree(VENDOR_DIR, ignore_errors=True)
        shutil.move(src, VENDOR_DIR)
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def _import_sefi():
    """Ensure the sefi package is importable, fetching source on first call."""
    if not os.path.isdir(os.path.join(VENDOR_DIR, "sefi")):
        _fetch_sefi_source()
    if VENDOR_DIR not in sys.path:
        sys.path.insert(0, VENDOR_DIR)
    from sefi import SEFIInferencePipeline  # noqa: WPS433

    return SEFIInferencePipeline


class SeFiLoader:
    """Load a SeFi-Image inference pipeline (downloads weights from HF on first use)."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "checkpoint": (CHECKPOINTS,),
                "dtype": (["auto", "bfloat16", "float16", "float32"],),
                "device": (["auto", "cuda", "cpu"],),
            }
        }

    RETURN_TYPES = ("SEFI_PIPE",)
    RETURN_NAMES = ("sefi_pipe",)
    FUNCTION = "load"
    CATEGORY = "SeFi-Image"

    # Cache one pipeline across executions; reload only when the key changes.
    _cache = {}

    def load(self, checkpoint, dtype, device):
        key = (checkpoint, dtype, device)
        cached = SeFiLoader._cache.get(key)
        if cached is not None:
            return (cached,)

        SEFIInferencePipeline = _import_sefi()

        kwargs = {}
        if dtype != "auto":
            # Their from_pretrained takes dtype as a string; pass the torch name.
            kwargs["dtype"] = dtype
        if device != "auto":
            kwargs["device"] = device

        pipe = SEFIInferencePipeline.from_pretrained(checkpoint, **kwargs)
        SeFiLoader._cache.clear()  # keep only one big model resident
        SeFiLoader._cache[key] = pipe
        return (pipe,)


class SeFiSampler:
    """Generate an image from a text prompt with a loaded SeFi-Image pipeline."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "sefi_pipe": ("SEFI_PIPE",),
                "prompt": ("STRING", {"multiline": True, "default": "A red apple on a wooden table."}),
                "steps": ("INT", {"default": 4, "min": 1, "max": 100}),
                "guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
                "width": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}),
                "height": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "SeFi-Image"

    def generate(self, sefi_pipe, prompt, steps, guidance_scale, width, height, seed):
        images = sefi_pipe(
            prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            seed=seed,
        )
        # sefi returns list[PIL.Image]; convert to a ComfyUI IMAGE batch [B,H,W,C].
        tensors = []
        for img in images:
            arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
            tensors.append(torch.from_numpy(arr))
        batch = torch.stack(tensors, dim=0)
        return (batch,)


NODE_CLASS_MAPPINGS = {
    "SeFiLoader": SeFiLoader,
    "SeFiSampler": SeFiSampler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SeFiLoader": "SeFi-Image Loader",
    "SeFiSampler": "SeFi-Image Sampler",
}