TashaSkyUp-EternalKernelPytorchNodes/hf_diffusers_nodes.py

import torch
import folder_paths
import torchvision
import torch.nn as nn
import functools
import torch
import json

try:
    from .config import config_settings
except ImportError as e:
    from config import config_settings

NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}


def ETK_HF_Diffusers_base(cls):
    cls.CATEGORY = "ETK/HF/Diffusers"
    # Add spaces to the camel case class name
    pretty_name = cls.__name__
    for i in range(1, len(pretty_name)):
        if pretty_name[i].isupper():
            pretty_name = pretty_name[:i] + " " + pretty_name[i:]
    cls.DISPLAY_NAME = pretty_name
    NODE_CLASS_MAPPINGS[cls.__name__] = cls
    NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name

    # Wrap the function defined in the FUNCTION attribute
    func_name = getattr(cls, "FUNCTION", None)
    if func_name and hasattr(cls, func_name):
        original_func = getattr(cls, func_name)

        @functools.wraps(original_func)
        def wrapped_func(*args, **kwargs):
            with torch.inference_mode(False):
                return original_func(*args, **kwargs)

        setattr(cls, func_name, wrapped_func)

    return cls


@ETK_HF_Diffusers_base
class DDPMPipline:
    """A node that loads the specified DDPMPipeline, supports only from_pretrained for now"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": ("STRING", {"default": "johnowhitaker/ddpm-butterflies-32px"}),
                "device": (["cpu", "cuda"], {"default": "cuda"}),
            },
        }

    RETURN_TYPES = ("DDPM_PL",)
    FUNCTION = "load_from_pretrained"

    def load_from_pretrained(self, model_name, device):
        from diffusers import DDPMPipeline
        ret = DDPMPipeline.from_pretrained(model_name, device=device)
        return (ret,)


@ETK_HF_Diffusers_base
class SampleDDPMPipline:
    """A node that samples from the specified DDPMPipeline"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("DDPM_PL", {}),
                "batch_size": ("INT", {"default": 1}),
                "num_timesteps": ("INT", {"default": 1000}),
            },
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "sample"

    def sample(
            self,
            model,
            batch_size,
            num_timesteps,
    ):
        ret = model(
            batch_size=batch_size,
            generator=None,
            num_inference_steps=num_timesteps,
            output_type="tensor",
        )
        ret = torch.tensor(ret.images)
        return (ret,)


@ETK_HF_Diffusers_base
class TensorToImage:
    """A node that converts a tensor to an image"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "tensor": ("TORCH_TENSOR", {}),
                "incoming_format": (["BHWC", "BCHW", "HWC"], {"default": "BHWC"}
                                    ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "tensor_to_image"

    def tensor_to_image(self, tensor, incoming_format):
        # outgoing format is BHWC torch tensor float 32 0-1

        if incoming_format == "BCHW":
            tensor = tensor.permute(0, 2, 3, 1)
        elif incoming_format == "HWC":
            tensor = tensor.unsqueeze(0)
        tensor = tensor.clamp(0, 1)

        return (tensor,)


@ETK_HF_Diffusers_base
class LoadHFDataset:
    """A node that loads the specified HFDataset"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dataset_name": ("STRING", {"default": "huggan/smithsonian_butterflies_subset"}),
            },
        }

    RETURN_TYPES = ("HF_DATASET",)
    FUNCTION = "load_from_hf"

    def load_from_hf(self, dataset_name):
        from datasets import load_dataset
        ret = load_dataset(dataset_name,
                           split="train",
                           )
        return (ret,)


@ETK_HF_Diffusers_base
class CreateUnet2dModel:
    """A node that creates a unet2d model"""

    @classmethod
    def INPUT_TYPES(cls):
        down_blocks = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")
        mid_blocks = ('UNetMidBlock2D', 'UnCLIPUNetMidBlock2D')
        up_blocks = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")

        return {
            "required": {
                "name": ("STRING", {"default": "UNet2D"}),
            },
            "optional": {
                "sample_size": ("TUPLE", {"default": (32, 32, 3,)}),
                "in_channels": ("INT", {"default": 3}),
                "out_channels": ("INT", {"default": 3}),
                "norm_num_groups": ("INT", {"default": 32}),
                "layers_per_block": ("INT", {"default": 2}),
                "down_block_types": ("TUPLE", {"default": (down_blocks[0],)}),
                "mid_block_scale_factor": ("FLOAT", {"default": 1.0}),
                "up_block_types": ("TUPLE", {"default": (up_blocks[0],)}),
                "block_out_channels": ("TUPLE", {"default": (64, 128, 256, 512,)}),
            },
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "create_unet2d_model"

    def create_unet2d_model(self, name, sample_size, in_channels, out_channels, norm_num_groups, layers_per_block,
                            down_block_types=("AttnDownBlock2D",),
                            mid_block_scale_factor=("UNetMidBlock2D",),
                            up_block_types=("AttnUpBlock2D",),
                            block_out_channels=(64, 128, 256, 512,)
                            ):
        from diffusers import UNet2DModel
        with torch.inference_mode(False):
            ret = UNet2DModel(
                sample_size=sample_size,
                in_channels=in_channels,
                out_channels=out_channels,
                norm_num_groups=norm_num_groups,
                layers_per_block=layers_per_block,
                down_block_types=down_block_types,
                mid_block_scale_factor=mid_block_scale_factor,
                up_block_types=up_block_types,
                block_out_channels=block_out_channels,

            )
        for param in ret.parameters():
            param.requires_grad = True
            #param.retains_grad = True

        return (ret,)

@ETK_HF_Diffusers_base
class TrainUnet2dModel:
    """A node that trains a unet2d model"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL", {}),
                "dataset": ("HF_DATASET", {}),
                "batch_size": ("INT", {"default": 1}),
                "num_epochs": ("INT", {"default": 1}),
                "learning_rate": ("FLOAT", {"default": 1e-4}),
            },
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "train_unet2d_model"

    def train_unet2d_model(self, model, dataset, batch_size, num_epochs, learning_rate):
        from torch.utils.data import DataLoader
        # import mse loss
        from torch.nn import MSELoss
        from torch.optim import Adam
        from diffusers import UNet2DModel


        model = model.train()
        model = model.cuda()

        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        optimizer = Adam(model.parameters(), lr=learning_rate)
        loss_fn = MSELoss()

        for epoch in range(num_epochs):
            for batch in dataloader:
                batch = batch["image"]
                batch = batch.cuda()
                optimizer.zero_grad()
                loss = loss_fn(model, batch)
                loss.backward()
                optimizer.step()
                print(loss.item())

        return (model,)