648 lines
23 KiB
Python
648 lines
23 KiB
Python
"""Runtime wrapper for local single-GPU DiffiT inference inside ComfyUI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from .utils import (
|
|
LOGGER,
|
|
default_cfg_scale_for_image_size,
|
|
get_hf_cache_dir,
|
|
get_preferred_diffit_download_dir,
|
|
infer_image_size_from_state_dict,
|
|
latent_size_for_image,
|
|
load_state_dict_file,
|
|
make_generator,
|
|
pick_device,
|
|
resolve_diffit_checkpoint_path,
|
|
to_comfy_image,
|
|
validate_image_size,
|
|
)
|
|
from .vendor.diffit_upstream import create_diffusion, diffusion_defaults
|
|
from .vendor.diffit_upstream.script_util import NUM_CLASSES
|
|
|
|
|
|
@dataclass
|
|
class DiffiTModelBundle:
|
|
checkpoint_name: str
|
|
checkpoint_path: str
|
|
image_size: int
|
|
latent_size: int
|
|
default_cfg_scale: float
|
|
model: torch.nn.Module
|
|
diffusion_cache: dict[int, Any] = field(default_factory=dict)
|
|
current_device: str = "cpu"
|
|
precision: str = "fp32"
|
|
|
|
|
|
@dataclass
|
|
class DiffiTConditioning:
|
|
class_id: int
|
|
label: str
|
|
|
|
|
|
@dataclass
|
|
class DiffiTLatent:
|
|
tensor: torch.Tensor
|
|
image_size: int
|
|
batch_size: int
|
|
seed: int
|
|
|
|
|
|
_MODEL_CACHE: dict[tuple[str, int], DiffiTModelBundle] = {}
|
|
_VAE_CACHE: dict[tuple[str, str], Any] = {}
|
|
|
|
|
|
def _lazy_import_diffit_model():
|
|
try:
|
|
module = importlib.import_module(".vendor.diffit_upstream.diffit", package=__package__)
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"DiffiT model import failed. Install this node pack's requirements first "
|
|
"(missing dependency is usually timm)."
|
|
) from exc
|
|
return module
|
|
|
|
|
|
def _lazy_import_autoencoder():
|
|
try:
|
|
module = importlib.import_module("diffusers.models")
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"DiffiT VAE decode requires diffusers. Install this node pack's requirements first."
|
|
) from exc
|
|
return module.AutoencoderKL
|
|
|
|
|
|
def _move_module(module: torch.nn.Module, device: torch.device) -> torch.nn.Module:
|
|
module.to(device)
|
|
module.eval()
|
|
return module
|
|
|
|
|
|
def _get_module_dtype(module: torch.nn.Module) -> torch.dtype:
|
|
try:
|
|
return next(module.parameters()).dtype
|
|
except StopIteration:
|
|
return torch.float32
|
|
|
|
|
|
def _make_progress_bar(total_steps: int, node_id=None):
|
|
try:
|
|
import comfy.utils # type: ignore
|
|
|
|
return comfy.utils.ProgressBar(total_steps, node_id=str(node_id) if node_id is not None else None)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _ensure_bundle_device(bundle: DiffiTModelBundle, device: torch.device) -> DiffiTModelBundle:
|
|
if bundle.current_device != str(device):
|
|
LOGGER.info("Moving DiffiT model to %s", device)
|
|
_move_module(bundle.model, device)
|
|
bundle.current_device = str(device)
|
|
return bundle
|
|
|
|
|
|
def _ensure_bundle_precision(bundle: DiffiTModelBundle, use_fp16: bool) -> DiffiTModelBundle:
|
|
target_precision = "fp16" if use_fp16 else "fp32"
|
|
target_dtype = torch.float16 if use_fp16 else torch.float32
|
|
current_dtype = _get_module_dtype(bundle.model)
|
|
if current_dtype != target_dtype:
|
|
LOGGER.info("Switching DiffiT model precision to %s", target_precision)
|
|
if use_fp16 and hasattr(bundle.model, "convert_to_fp16"):
|
|
bundle.model.convert_to_fp16()
|
|
elif not use_fp16 and hasattr(bundle.model, "convert_to_fp32"):
|
|
bundle.model.convert_to_fp32()
|
|
else:
|
|
bundle.model.to(dtype=target_dtype)
|
|
bundle.precision = target_precision
|
|
bundle.model.eval()
|
|
return bundle
|
|
|
|
|
|
def _build_diffusion(steps: int):
|
|
config = diffusion_defaults()
|
|
config["timestep_respacing"] = str(int(steps))
|
|
return create_diffusion(**config)
|
|
|
|
|
|
def _get_cached_diffusion(bundle: DiffiTModelBundle, steps: int):
|
|
if steps not in bundle.diffusion_cache:
|
|
bundle.diffusion_cache[steps] = _build_diffusion(steps)
|
|
return bundle.diffusion_cache[steps]
|
|
|
|
|
|
def _resolve_cfg_eps_channels(model_bundle: DiffiTModelBundle, cfg_channel_mode: str) -> int:
|
|
if cfg_channel_mode == "latent_channels":
|
|
return int(getattr(model_bundle.model, "in_channels", 4))
|
|
if cfg_channel_mode == "upstream_repo_3ch":
|
|
return 3
|
|
raise ValueError(
|
|
f"Unsupported cfg_channel_mode '{cfg_channel_mode}'. "
|
|
"Use 'latent_channels' or 'upstream_repo_3ch'."
|
|
)
|
|
|
|
|
|
def _download_checkpoint_from_hf(checkpoint_name: str) -> Path:
|
|
if checkpoint_name not in {"diffit_256.safetensors", "diffit_512.safetensors"}:
|
|
raise FileNotFoundError(
|
|
f"Checkpoint {checkpoint_name} was not found locally and is not one of the official DiffiT filenames."
|
|
)
|
|
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
from huggingface_hub.utils import HfHubHTTPError
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"Downloading DiffiT weights requires huggingface_hub. Install this node pack's requirements first."
|
|
) from exc
|
|
|
|
target_dir = get_preferred_diffit_download_dir()
|
|
LOGGER.info("Downloading %s into %s", checkpoint_name, target_dir)
|
|
try:
|
|
local_path = hf_hub_download(
|
|
repo_id="nvidia/DiffiT",
|
|
filename=checkpoint_name,
|
|
local_dir=str(target_dir),
|
|
)
|
|
except HfHubHTTPError as exc:
|
|
raise RuntimeError(
|
|
"Failed to download the official DiffiT checkpoint from Hugging Face. "
|
|
"Check your network connection and any Hugging Face auth configuration."
|
|
) from exc
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Unexpected DiffiT download failure for {checkpoint_name}.") from exc
|
|
|
|
return Path(local_path)
|
|
|
|
|
|
def load_diffit_model(
|
|
checkpoint_name: str,
|
|
checkpoint_path: str = "",
|
|
image_size_override: str = "auto",
|
|
download_if_missing: bool = True,
|
|
device: str = "auto",
|
|
) -> DiffiTModelBundle:
|
|
resolved_path = resolve_diffit_checkpoint_path(checkpoint_name, checkpoint_path)
|
|
if resolved_path is None and download_if_missing:
|
|
resolved_path = _download_checkpoint_from_hf(checkpoint_name)
|
|
|
|
if resolved_path is None:
|
|
target_dir = get_preferred_diffit_download_dir()
|
|
raise FileNotFoundError(
|
|
f"DiffiT checkpoint {checkpoint_name} was not found. Place it in {target_dir} "
|
|
f"or in {target_dir.parent.parent / 'Packages' / 'ComfyUI-rect' / 'models' / 'diffusion_models'}."
|
|
)
|
|
|
|
state_dict = load_state_dict_file(resolved_path)
|
|
inferred_size = infer_image_size_from_state_dict(state_dict)
|
|
image_size = inferred_size if image_size_override == "auto" else validate_image_size(int(image_size_override))
|
|
cache_key = (str(resolved_path.resolve()), image_size)
|
|
|
|
bundle = _MODEL_CACHE.get(cache_key)
|
|
if bundle is None:
|
|
diffit_module = _lazy_import_diffit_model()
|
|
latent_size = latent_size_for_image(image_size)
|
|
model = diffit_module.Diffit(input_size=latent_size)
|
|
load_result = model.load_state_dict(state_dict, strict=True)
|
|
if getattr(load_result, "missing_keys", None) or getattr(load_result, "unexpected_keys", None):
|
|
LOGGER.warning(
|
|
"DiffiT checkpoint load reported mismatched keys. Missing=%s Unexpected=%s",
|
|
getattr(load_result, "missing_keys", []),
|
|
getattr(load_result, "unexpected_keys", []),
|
|
)
|
|
model.eval()
|
|
bundle = DiffiTModelBundle(
|
|
checkpoint_name=resolved_path.name,
|
|
checkpoint_path=str(resolved_path),
|
|
image_size=image_size,
|
|
latent_size=latent_size,
|
|
default_cfg_scale=default_cfg_scale_for_image_size(image_size),
|
|
model=model,
|
|
diffusion_cache={250: _build_diffusion(250)},
|
|
current_device="cpu",
|
|
)
|
|
_MODEL_CACHE[cache_key] = bundle
|
|
|
|
target_device = pick_device(device)
|
|
return _ensure_bundle_device(bundle, target_device)
|
|
|
|
|
|
def create_class_conditioning(
|
|
class_id: int,
|
|
label_choice: str = "<use class_id>",
|
|
) -> DiffiTConditioning:
|
|
resolved_class_id = int(class_id)
|
|
resolved_label = f"class_{resolved_class_id}"
|
|
|
|
if label_choice and label_choice != "<use class_id>":
|
|
prefix = label_choice.split(":", 1)[0].strip()
|
|
try:
|
|
resolved_class_id = int(prefix)
|
|
except ValueError:
|
|
raise ValueError(f"Unable to parse ImageNet class id from label choice: {label_choice}")
|
|
resolved_label = label_choice.split(":", 1)[1].strip() if ":" in label_choice else label_choice
|
|
|
|
if not 0 <= resolved_class_id < NUM_CLASSES:
|
|
raise ValueError(f"DiffiT class id must be between 0 and 999. Received {resolved_class_id}.")
|
|
|
|
return DiffiTConditioning(class_id=resolved_class_id, label=resolved_label)
|
|
|
|
|
|
def create_empty_latent(batch_size: int, image_size: int, seed: int) -> DiffiTLatent:
|
|
image_size = validate_image_size(int(image_size))
|
|
batch_size = int(batch_size)
|
|
if batch_size < 1:
|
|
raise ValueError("DiffiT batch_size must be at least 1.")
|
|
|
|
latent_size = latent_size_for_image(image_size)
|
|
generator = torch.Generator(device="cpu")
|
|
generator.manual_seed(int(seed) & 0xFFFFFFFFFFFFFFFF)
|
|
tensor = torch.randn((batch_size, 4, latent_size, latent_size), generator=generator, dtype=torch.float32)
|
|
return DiffiTLatent(tensor=tensor, image_size=image_size, batch_size=batch_size, seed=int(seed))
|
|
|
|
|
|
def sample_diffit(
|
|
model_bundle: DiffiTModelBundle,
|
|
conditioning: DiffiTConditioning,
|
|
latent: DiffiTLatent,
|
|
steps: int = 250,
|
|
cfg_scale: float = 0.0,
|
|
seed: int = 0,
|
|
use_ddim: bool = False,
|
|
use_fp16: bool = True,
|
|
scale_pow: float = 4.0,
|
|
cfg_enabled: bool = True,
|
|
cfg_channel_mode: str = "latent_channels",
|
|
node_id=None,
|
|
) -> DiffiTLatent:
|
|
if int(steps) < 2:
|
|
raise ValueError("DiffiT sampling requires at least 2 steps.")
|
|
|
|
if latent.image_size != model_bundle.image_size:
|
|
raise ValueError(
|
|
f"Latent image_size={latent.image_size} does not match loaded DiffiT model image_size={model_bundle.image_size}."
|
|
)
|
|
|
|
device = torch.device(model_bundle.current_device)
|
|
bundle = _ensure_bundle_device(model_bundle, device)
|
|
fp16_requested = bool(use_fp16 and device.type == "cuda")
|
|
_ensure_bundle_precision(bundle, fp16_requested)
|
|
diffusion = _get_cached_diffusion(bundle, int(steps))
|
|
latent_tensor = latent.tensor.to(device=device, dtype=torch.float32, non_blocking=device.type == "cuda")
|
|
|
|
model_kwargs: dict[str, Any] = {}
|
|
sample_model = bundle.model.forward
|
|
if cfg_enabled:
|
|
z = torch.cat([latent_tensor, latent_tensor], dim=0)
|
|
classes = torch.full((latent.batch_size,), conditioning.class_id, device=device, dtype=torch.long)
|
|
classes_null = torch.full((latent.batch_size,), NUM_CLASSES, device=device, dtype=torch.long)
|
|
model_kwargs["y"] = torch.cat([classes, classes_null], dim=0)
|
|
model_kwargs["cfg_scale"] = float(cfg_scale) if float(cfg_scale) > 0 else bundle.default_cfg_scale
|
|
model_kwargs["diffusion_steps"] = diffusion_defaults()["diffusion_steps"]
|
|
model_kwargs["scale_pow"] = float(scale_pow)
|
|
model_kwargs["cfg_eps_channels"] = _resolve_cfg_eps_channels(bundle, cfg_channel_mode)
|
|
sample_model = bundle.model.forward_with_cfg
|
|
else:
|
|
z = latent_tensor
|
|
model_kwargs["y"] = torch.full(
|
|
(latent.batch_size,),
|
|
conditioning.class_id,
|
|
device=device,
|
|
dtype=torch.long,
|
|
)
|
|
|
|
sample_fn = diffusion.ddim_sample_loop if use_ddim else diffusion.p_sample_loop
|
|
progress_bar = _make_progress_bar(int(steps), node_id=node_id)
|
|
|
|
if device.type == "cuda":
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
def run_sampler(autocast_enabled: bool) -> torch.Tensor:
|
|
generator = make_generator(device, int(seed))
|
|
step_callback = None
|
|
if progress_bar is not None:
|
|
step_callback = lambda current, total: progress_bar.update_absolute(current, total=total)
|
|
with torch.inference_mode():
|
|
with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=autocast_enabled):
|
|
return sample_fn(
|
|
sample_model,
|
|
z.shape,
|
|
noise=z,
|
|
clip_denoised=False,
|
|
progress=False,
|
|
model_kwargs=model_kwargs,
|
|
device=device,
|
|
generator=generator,
|
|
step_callback=step_callback,
|
|
)
|
|
|
|
def run_sampler_with_fallback() -> torch.Tensor:
|
|
autocast_enabled = fp16_requested
|
|
try:
|
|
sampled_tensor = run_sampler(autocast_enabled)
|
|
except RuntimeError as exc:
|
|
if not fp16_requested:
|
|
raise
|
|
LOGGER.warning("DiffiT fp16 sampler failed (%s). Retrying in fp32.", exc)
|
|
_ensure_bundle_precision(bundle, False)
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(0, total=int(steps))
|
|
return run_sampler(False)
|
|
|
|
if fp16_requested and not torch.isfinite(sampled_tensor).all():
|
|
LOGGER.warning("DiffiT fp16 sampler produced non-finite values. Retrying in fp32.")
|
|
_ensure_bundle_precision(bundle, False)
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(0, total=int(steps))
|
|
return run_sampler(False)
|
|
return sampled_tensor
|
|
|
|
sampled = run_sampler_with_fallback()
|
|
|
|
if cfg_enabled:
|
|
sampled, _ = sampled.chunk(2, dim=0)
|
|
|
|
if not torch.isfinite(sampled).all():
|
|
raise RuntimeError("DiffiT sampler produced non-finite latent values.")
|
|
|
|
return DiffiTLatent(
|
|
tensor=sampled.detach().cpu().float(),
|
|
image_size=bundle.image_size,
|
|
batch_size=latent.batch_size,
|
|
seed=int(seed),
|
|
)
|
|
|
|
|
|
def sample_diffit_official(
|
|
model_bundle: DiffiTModelBundle,
|
|
conditioning: DiffiTConditioning,
|
|
latent: DiffiTLatent,
|
|
steps: int = 250,
|
|
cfg_scale: float = 0.0,
|
|
use_ddim: bool = False,
|
|
use_fp16: bool = False,
|
|
scale_pow: float = 4.0,
|
|
cfg_cond: bool = True,
|
|
clip_denoised: bool = False,
|
|
node_id=None,
|
|
) -> DiffiTLatent:
|
|
"""
|
|
Follow the released NVlabs sample.py path as closely as practical in ComfyUI.
|
|
|
|
This intentionally preserves the public repo's CFG behavior by applying
|
|
guidance through ``forward_with_cfg`` with 3 guided channels.
|
|
"""
|
|
if int(steps) < 2:
|
|
raise ValueError("DiffiT official sampling requires at least 2 steps.")
|
|
|
|
if latent.image_size != model_bundle.image_size:
|
|
raise ValueError(
|
|
f"Latent image_size={latent.image_size} does not match loaded DiffiT model image_size={model_bundle.image_size}."
|
|
)
|
|
|
|
device = torch.device(model_bundle.current_device)
|
|
if use_fp16 and device.type != "cuda":
|
|
raise RuntimeError("DiffiT official fp16 mode requires a CUDA device.")
|
|
|
|
bundle = _ensure_bundle_device(model_bundle, device)
|
|
_ensure_bundle_precision(bundle, bool(use_fp16 and device.type == "cuda"))
|
|
diffusion = _get_cached_diffusion(bundle, int(steps))
|
|
|
|
model_dtype = _get_module_dtype(bundle.model)
|
|
latent_tensor = latent.tensor.to(device=device, dtype=model_dtype, non_blocking=device.type == "cuda")
|
|
|
|
model_kwargs: dict[str, Any] = {}
|
|
if cfg_cond:
|
|
z = torch.cat([latent_tensor, latent_tensor], dim=0)
|
|
classes = torch.full((latent.batch_size,), conditioning.class_id, device=device, dtype=torch.long)
|
|
classes_null = torch.full((latent.batch_size,), NUM_CLASSES, device=device, dtype=torch.long)
|
|
model_kwargs["y"] = torch.cat([classes, classes_null], dim=0)
|
|
model_kwargs["cfg_scale"] = float(cfg_scale) if float(cfg_scale) > 0 else bundle.default_cfg_scale
|
|
model_kwargs["diffusion_steps"] = diffusion_defaults()["diffusion_steps"]
|
|
model_kwargs["scale_pow"] = float(scale_pow)
|
|
else:
|
|
z = latent_tensor
|
|
model_kwargs["y"] = torch.full(
|
|
(latent.batch_size,),
|
|
conditioning.class_id,
|
|
device=device,
|
|
dtype=torch.long,
|
|
)
|
|
|
|
sample_fn = diffusion.ddim_sample_loop if use_ddim else diffusion.p_sample_loop
|
|
progress_bar = _make_progress_bar(int(steps), node_id=node_id)
|
|
|
|
if device.type == "cuda":
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
step_callback = None
|
|
if progress_bar is not None:
|
|
step_callback = lambda current, total: progress_bar.update_absolute(current, total=total)
|
|
|
|
with torch.inference_mode():
|
|
sampled = sample_fn(
|
|
bundle.model.forward_with_cfg,
|
|
z.shape,
|
|
z,
|
|
clip_denoised=bool(clip_denoised),
|
|
progress=False,
|
|
model_kwargs=model_kwargs,
|
|
device=device,
|
|
step_callback=step_callback,
|
|
)
|
|
|
|
if cfg_cond:
|
|
sampled, _ = sampled.chunk(2, dim=0)
|
|
|
|
if not torch.isfinite(sampled).all():
|
|
raise RuntimeError(
|
|
"DiffiT official sampler produced non-finite latent values. "
|
|
"This follows the released repo path closely; disable fp16 or use the non-official sampler if needed."
|
|
)
|
|
|
|
return DiffiTLatent(
|
|
tensor=sampled.detach().cpu().float(),
|
|
image_size=bundle.image_size,
|
|
batch_size=latent.batch_size,
|
|
seed=latent.seed,
|
|
)
|
|
|
|
|
|
def _load_vae(variant: str, device: torch.device):
|
|
variant = variant.lower()
|
|
if variant not in {"ema", "mse"}:
|
|
raise ValueError(f"Unsupported DiffiT VAE variant '{variant}'. Use 'ema' or 'mse'.")
|
|
|
|
cache_key = (variant, str(device))
|
|
cached = _VAE_CACHE.get(cache_key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
AutoencoderKL = _lazy_import_autoencoder()
|
|
repo_id = f"stabilityai/sd-vae-ft-{variant}"
|
|
cache_dir = get_hf_cache_dir(f"sd-vae-ft-{variant}")
|
|
try:
|
|
vae = AutoencoderKL.from_pretrained(repo_id, cache_dir=str(cache_dir))
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
"Failed to load the DiffiT VAE from Hugging Face. "
|
|
"Check your network connection, local Hugging Face cache, and any auth configuration."
|
|
) from exc
|
|
|
|
vae = _move_module(vae, device)
|
|
if hasattr(vae, "enable_slicing"):
|
|
vae.enable_slicing()
|
|
_VAE_CACHE[cache_key] = vae
|
|
return vae
|
|
|
|
|
|
def decode_diffit_latent(
|
|
latent: DiffiTLatent,
|
|
vae_variant: str = "ema",
|
|
use_fp16: bool = True,
|
|
device: str = "auto",
|
|
) -> torch.Tensor:
|
|
target_device = pick_device(device)
|
|
vae = _load_vae(vae_variant, target_device)
|
|
decode_fp16 = bool(use_fp16 and target_device.type == "cuda")
|
|
vae.to(dtype=torch.float16 if decode_fp16 else torch.float32)
|
|
latent_tensor = latent.tensor.to(target_device, non_blocking=target_device.type == "cuda")
|
|
|
|
def run_decode(autocast_enabled: bool) -> torch.Tensor:
|
|
with torch.inference_mode():
|
|
with torch.autocast(
|
|
device_type=target_device.type,
|
|
dtype=torch.float16,
|
|
enabled=autocast_enabled,
|
|
):
|
|
return vae.decode(latent_tensor / 0.18215).sample
|
|
|
|
def run_decode_with_fallback() -> torch.Tensor:
|
|
try:
|
|
decoded_tensor = run_decode(decode_fp16)
|
|
except RuntimeError as exc:
|
|
if not decode_fp16:
|
|
raise
|
|
LOGGER.warning("DiffiT fp16 VAE decode failed (%s). Retrying in fp32.", exc)
|
|
vae.to(dtype=torch.float32)
|
|
return run_decode(False)
|
|
|
|
if decode_fp16 and not torch.isfinite(decoded_tensor).all():
|
|
LOGGER.warning("DiffiT fp16 VAE decode produced non-finite values. Retrying in fp32.")
|
|
vae.to(dtype=torch.float32)
|
|
return run_decode(False)
|
|
return decoded_tensor
|
|
|
|
decoded = run_decode_with_fallback()
|
|
|
|
if not torch.isfinite(decoded).all():
|
|
raise RuntimeError("DiffiT VAE decode produced non-finite image values.")
|
|
|
|
return to_comfy_image(decoded)
|
|
|
|
|
|
def decode_diffit_latent_official(
|
|
latent: DiffiTLatent,
|
|
vae_variant: str = "ema",
|
|
device: str = "auto",
|
|
) -> torch.Tensor:
|
|
"""
|
|
Follow the released sample.py VAE path: decode in the default VAE dtype.
|
|
"""
|
|
target_device = pick_device(device)
|
|
vae = _load_vae(vae_variant, target_device)
|
|
vae.to(dtype=torch.float32)
|
|
latent_tensor = latent.tensor.to(
|
|
target_device,
|
|
dtype=torch.float32,
|
|
non_blocking=target_device.type == "cuda",
|
|
)
|
|
|
|
with torch.inference_mode():
|
|
decoded = vae.decode(latent_tensor / 0.18215).sample
|
|
|
|
if not torch.isfinite(decoded).all():
|
|
raise RuntimeError("DiffiT official VAE decode produced non-finite image values.")
|
|
|
|
return to_comfy_image(decoded)
|
|
|
|
|
|
def generate_diffit_image(
|
|
model_bundle: DiffiTModelBundle,
|
|
conditioning: DiffiTConditioning,
|
|
batch_size: int,
|
|
latent_seed: int,
|
|
steps: int = 250,
|
|
cfg_scale: float = 0.0,
|
|
sampler_seed: int = 0,
|
|
use_ddim: bool = False,
|
|
use_fp16: bool = True,
|
|
scale_pow: float = 4.0,
|
|
cfg_enabled: bool = True,
|
|
cfg_channel_mode: str = "latent_channels",
|
|
vae_variant: str = "ema",
|
|
node_id=None,
|
|
) -> tuple[torch.Tensor, DiffiTLatent]:
|
|
latent = create_empty_latent(batch_size=batch_size, image_size=model_bundle.image_size, seed=latent_seed)
|
|
sampled = sample_diffit(
|
|
model_bundle=model_bundle,
|
|
conditioning=conditioning,
|
|
latent=latent,
|
|
steps=steps,
|
|
cfg_scale=cfg_scale,
|
|
seed=sampler_seed,
|
|
use_ddim=use_ddim,
|
|
use_fp16=use_fp16,
|
|
scale_pow=scale_pow,
|
|
cfg_enabled=cfg_enabled,
|
|
cfg_channel_mode=cfg_channel_mode,
|
|
node_id=node_id,
|
|
)
|
|
image = decode_diffit_latent(sampled, vae_variant=vae_variant, use_fp16=use_fp16, device=model_bundle.current_device)
|
|
return image, sampled
|
|
|
|
|
|
def generate_diffit_image_official(
|
|
model_bundle: DiffiTModelBundle,
|
|
conditioning: DiffiTConditioning,
|
|
batch_size: int,
|
|
latent_seed: int,
|
|
steps: int = 250,
|
|
cfg_scale: float = 0.0,
|
|
use_ddim: bool = False,
|
|
use_fp16: bool = False,
|
|
scale_pow: float = 4.0,
|
|
cfg_cond: bool = True,
|
|
clip_denoised: bool = False,
|
|
vae_variant: str = "ema",
|
|
node_id=None,
|
|
) -> tuple[torch.Tensor, DiffiTLatent]:
|
|
latent = create_empty_latent(batch_size=batch_size, image_size=model_bundle.image_size, seed=latent_seed)
|
|
sampled = sample_diffit_official(
|
|
model_bundle=model_bundle,
|
|
conditioning=conditioning,
|
|
latent=latent,
|
|
steps=steps,
|
|
cfg_scale=cfg_scale,
|
|
use_ddim=use_ddim,
|
|
use_fp16=use_fp16,
|
|
scale_pow=scale_pow,
|
|
cfg_cond=cfg_cond,
|
|
clip_denoised=clip_denoised,
|
|
node_id=node_id,
|
|
)
|
|
image = decode_diffit_latent_official(
|
|
sampled,
|
|
vae_variant=vae_variant,
|
|
device=model_bundle.current_device,
|
|
)
|
|
return image, sampled
|