97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
"""ComfyUI node definitions for the official single-node DiffiT pipeline."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .diffit_runtime import (
|
|
create_class_conditioning,
|
|
generate_diffit_image_official,
|
|
load_diffit_model,
|
|
)
|
|
from .utils import LOGGER, list_available_diffit_checkpoints
|
|
|
|
|
|
CATEGORY = "NVIDIA/DiffiT"
|
|
|
|
|
|
class DiffiTOfficialImageNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"checkpoint_name": (list_available_diffit_checkpoints(),),
|
|
"download_if_missing": ("BOOLEAN", {"default": True}),
|
|
"device": (["auto", "cuda", "cpu"],),
|
|
"class_id": ("INT", {"default": 207, "min": 0, "max": 999}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
|
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7FFFFFFFFFFFFFFF}),
|
|
"steps": ("INT", {"default": 250, "min": 2, "max": 1000}),
|
|
"cfg_scale": ("FLOAT", {"default": 4.4, "min": 0.0, "max": 32.0, "step": 0.01}),
|
|
"cfg_cond": ("BOOLEAN", {"default": True}),
|
|
"use_ddim": ("BOOLEAN", {"default": False}),
|
|
"scale_pow": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 16.0, "step": 0.1}),
|
|
"clip_denoised": ("BOOLEAN", {"default": False}),
|
|
"vae_variant": (["ema", "mse"],),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = CATEGORY
|
|
|
|
def generate(
|
|
self,
|
|
checkpoint_name,
|
|
download_if_missing,
|
|
device,
|
|
class_id,
|
|
batch_size,
|
|
seed,
|
|
steps,
|
|
cfg_scale,
|
|
cfg_cond,
|
|
use_ddim,
|
|
scale_pow,
|
|
clip_denoised,
|
|
vae_variant,
|
|
unique_id=None,
|
|
):
|
|
try:
|
|
model = load_diffit_model(
|
|
checkpoint_name=checkpoint_name,
|
|
checkpoint_path="",
|
|
image_size_override="auto",
|
|
download_if_missing=download_if_missing,
|
|
device=device,
|
|
)
|
|
conditioning = create_class_conditioning(class_id, "<use class_id>")
|
|
image, _sampled = generate_diffit_image_official(
|
|
model_bundle=model,
|
|
conditioning=conditioning,
|
|
batch_size=batch_size,
|
|
latent_seed=seed,
|
|
steps=steps,
|
|
cfg_scale=cfg_scale,
|
|
cfg_cond=cfg_cond,
|
|
use_ddim=use_ddim,
|
|
use_fp16=False,
|
|
scale_pow=scale_pow,
|
|
clip_denoised=clip_denoised,
|
|
vae_variant=vae_variant,
|
|
node_id=unique_id,
|
|
)
|
|
return (image,)
|
|
except Exception as exc:
|
|
LOGGER.exception("DiffiT official image generation failed.")
|
|
raise RuntimeError(f"DiffiT Official Image failed: {exc}") from exc
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DiffiTOfficialImage": DiffiTOfficialImageNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DiffiTOfficialImage": "DiffiT Official Image",
|
|
}
|