117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
import os
|
|
import torch
|
|
import folder_paths
|
|
import comfy.model_management
|
|
import comfy.model_patcher
|
|
import comfy.utils
|
|
import comfy.latent_formats
|
|
|
|
from .models import DiT_models
|
|
from .diffusion import create_diffusion
|
|
|
|
# load these from separate folder
|
|
folder_paths.folder_names_and_paths["dit"] = (
|
|
[os.path.join(folder_paths.models_dir,"dit")],
|
|
folder_paths.supported_pt_extensions
|
|
)
|
|
|
|
class DiTCheckpointLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("dit"),),
|
|
"model": (list(DiT_models.keys()),),
|
|
"image_size": ([256, 512],),
|
|
"num_classes": ("INT", {"default": 1000, "min": 0,}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("DIT",) # could be MODEL if it is made compatible?
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "DiT"
|
|
TITLE = "DiTCheckpointLoader"
|
|
|
|
def load_checkpoint(self, ckpt_name, model, image_size, num_classes):
|
|
# note: switch to custom comfy.model_base eventually
|
|
model = DiT_models[model](
|
|
input_size=image_size // 8, # latent size
|
|
num_classes=num_classes
|
|
)
|
|
|
|
ckpt_path = folder_paths.get_full_path("dit", ckpt_name)
|
|
state_dict = comfy.utils.load_torch_file(ckpt_path)
|
|
model.load_state_dict(state_dict)
|
|
model.eval() # important, apparently
|
|
|
|
# need these later anyway
|
|
model.latent_format = comfy.latent_formats.SD15()
|
|
model.latent_size = image_size // 8
|
|
model.num_classes = num_classes
|
|
|
|
# I didn't expect this to work but it looks like it does.
|
|
model_patcher = comfy.model_patcher.ModelPatcher(
|
|
model,
|
|
load_device=comfy.model_management.get_torch_device(),
|
|
offload_device=comfy.model_management.unet_offload_device(),
|
|
current_device="cpu"
|
|
)
|
|
|
|
# return (model,)
|
|
return (model_patcher,)
|
|
|
|
class DiTSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("DIT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
|
"class_labels": ([207,],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "DiT"
|
|
TITLE = "DiTSampler"
|
|
|
|
def sample(self, model, seed, steps, cfg, batch_size, class_labels):
|
|
device = comfy.model_management.get_torch_device()
|
|
diffusion = create_diffusion(str(steps))
|
|
|
|
# pre
|
|
comfy.model_management.load_model_gpu(model)
|
|
real_model = model.model
|
|
|
|
# Create sampling noise:
|
|
z = torch.randn(batch_size, 4, real_model.latent_size, real_model.latent_size, device=device)
|
|
y = torch.tensor([class_labels] * batch_size, device=device)
|
|
|
|
# Setup classifier-free guidance:
|
|
z = torch.cat([z, z], 0)
|
|
y_null = torch.tensor([1000] * batch_size, device=device)
|
|
y = torch.cat([y, y_null], 0)
|
|
model_kwargs = dict(y=y, cfg_scale=cfg)
|
|
|
|
# Sample images:
|
|
samples = diffusion.p_sample_loop(
|
|
model.model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
|
|
)
|
|
samples, _ = samples.chunk(2, dim=0) # Remove null class samples
|
|
samples = real_model.latent_format.process_out(samples.to(torch.float32))
|
|
samples = samples.cpu()
|
|
|
|
return ({"samples": samples},)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DiTCheckpointLoader": DiTCheckpointLoader,
|
|
"DiTSampler": DiTSampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DiTCheckpointLoader": DiTCheckpointLoader.TITLE,
|
|
"DiTSampler": DiTSampler.TITLE,
|
|
}
|