|
|
|
|
| import torch
|
| import numpy as np
|
| from typing import Tuple, TypedDict, Callable
|
|
|
| import comfy.model_management
|
| from comfy.sd import load_unet
|
| from comfy.ldm.models.autoencoder import AutoencoderKL
|
| from comfy.model_base import BaseModel
|
| from comfy.model_patcher import ModelPatcher
|
| from PIL import Image
|
| from nodes import VAEEncode
|
| from ...libs.image import np2tensor, pil2tensor
|
|
|
| class UnetParams(TypedDict):
|
| input: torch.Tensor
|
| timestep: torch.Tensor
|
| c: dict
|
| cond_or_uncond: torch.Tensor
|
|
|
| class VAEEncodeArgMax(VAEEncode):
|
| def encode(self, vae, pixels):
|
| assert isinstance(
|
| vae.first_stage_model, AutoencoderKL
|
| ), "ArgMax only supported for AutoencoderKL"
|
| original_sample_mode = vae.first_stage_model.regularization.sample
|
| vae.first_stage_model.regularization.sample = False
|
| ret = super().encode(vae, pixels)
|
| vae.first_stage_model.regularization.sample = original_sample_mode
|
| return ret
|
|
|
| class ICLight:
|
|
|
| @staticmethod
|
| def apply_c_concat(params: UnetParams, concat_conds) -> UnetParams:
|
| """Apply c_concat on unet call."""
|
| sample = params["input"]
|
| params["c"]["c_concat"] = torch.cat(
|
| (
|
| [concat_conds.to(sample.device)]
|
| * (sample.shape[0] // concat_conds.shape[0])
|
| ),
|
| dim=0,
|
| )
|
| return params
|
|
|
| @staticmethod
|
| def create_custom_conv(
|
| original_conv: torch.nn.Module,
|
| dtype: torch.dtype,
|
| device=torch.device,
|
| ) -> torch.nn.Module:
|
| with torch.no_grad():
|
| new_conv_in = torch.nn.Conv2d(
|
| 8,
|
| original_conv.out_channels,
|
| original_conv.kernel_size,
|
| original_conv.stride,
|
| original_conv.padding,
|
| )
|
| new_conv_in.weight.zero_()
|
| new_conv_in.weight[:, :4, :, :].copy_(original_conv.weight)
|
| new_conv_in.bias = original_conv.bias
|
| return new_conv_in.to(dtype=dtype, device=device)
|
|
|
| def generate_lighting_image(self, original_image, direction):
|
| _, image_height, image_width, _ = original_image.shape
|
| if direction == 'Left Light':
|
| gradient = np.linspace(255, 0, image_width)
|
| image = np.tile(gradient, (image_height, 1))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif direction == 'Right Light':
|
| gradient = np.linspace(0, 255, image_width)
|
| image = np.tile(gradient, (image_height, 1))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif direction == 'Top Light':
|
| gradient = np.linspace(255, 0, image_height)[:, None]
|
| image = np.tile(gradient, (1, image_width))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif direction == 'Bottom Light':
|
| gradient = np.linspace(0, 255, image_height)[:, None]
|
| image = np.tile(gradient, (1, image_width))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif direction == 'Circle Light':
|
| x = np.linspace(-1, 1, image_width)
|
| y = np.linspace(-1, 1, image_height)
|
| x, y = np.meshgrid(x, y)
|
| r = np.sqrt(x ** 2 + y ** 2)
|
| r = r / r.max()
|
| color1 = np.array([0, 0, 0])[np.newaxis, np.newaxis, :]
|
| color2 = np.array([255, 255, 255])[np.newaxis, np.newaxis, :]
|
| gradient = (color1 * r[..., np.newaxis] + color2 * (1 - r)[..., np.newaxis]).astype(np.uint8)
|
| image = pil2tensor(Image.fromarray(gradient))
|
| return image
|
| else:
|
| image = pil2tensor(Image.new('RGB', (1, 1), (0, 0, 0)))
|
| return image
|
|
|
| def generate_source_image(self, original_image, source):
|
| batch_size, image_height, image_width, _ = original_image.shape
|
| if source == 'Use Flipped Background Image':
|
| if batch_size < 2:
|
| raise ValueError('Must be at least 2 image to use flipped background image.')
|
| original_image = [img.unsqueeze(0) for img in original_image]
|
| image = torch.flip(original_image[1], [2])
|
| return image
|
| elif source == 'Ambient':
|
| input_bg = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8) + 64
|
| return np2tensor(input_bg)
|
| elif source == 'Left Light':
|
| gradient = np.linspace(224, 32, image_width)
|
| image = np.tile(gradient, (image_height, 1))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif source == 'Right Light':
|
| gradient = np.linspace(32, 224, image_width)
|
| image = np.tile(gradient, (image_height, 1))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif source == 'Top Light':
|
| gradient = np.linspace(224, 32, image_height)[:, None]
|
| image = np.tile(gradient, (1, image_width))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| elif source == 'Bottom Light':
|
| gradient = np.linspace(32, 224, image_height)[:, None]
|
| image = np.tile(gradient, (1, image_width))
|
| input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
|
| return np2tensor(input_bg)
|
| else:
|
| image = pil2tensor(Image.new('RGB', (1, 1), (0, 0, 0)))
|
| return image
|
|
|
|
|
| def apply(self, ic_model_path, model, c_concat: dict, ic_model=None) -> Tuple[ModelPatcher]:
|
| device = comfy.model_management.get_torch_device()
|
| dtype = comfy.model_management.unet_dtype()
|
| work_model = model.clone()
|
|
|
|
|
| base_model: BaseModel = work_model.model
|
| scale_factor = base_model.model_config.latent_format.scale_factor
|
|
|
|
|
| concat_conds: torch.Tensor = c_concat["samples"] * scale_factor
|
|
|
| concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)
|
|
|
| def unet_dummy_apply(unet_apply: Callable, params: UnetParams):
|
| """A dummy unet apply wrapper serving as the endpoint of wrapper
|
| chain."""
|
| return unet_apply(x=params["input"], t=params["timestep"], **params["c"])
|
|
|
| existing_wrapper = work_model.model_options.get(
|
| "model_function_wrapper", unet_dummy_apply
|
| )
|
|
|
| def wrapper_func(unet_apply: Callable, params: UnetParams):
|
| return existing_wrapper(unet_apply, params=self.apply_c_concat(params, concat_conds))
|
|
|
| work_model.set_model_unet_function_wrapper(wrapper_func)
|
| if not ic_model:
|
| ic_model = load_unet(ic_model_path)
|
| ic_model_state_dict = ic_model.model.diffusion_model.state_dict()
|
|
|
| work_model.add_patches(
|
| patches={
|
| ("diffusion_model." + key): (
|
| 'diff',
|
| [
|
| value.to(dtype=dtype, device=device),
|
| {"pad_weight": key == 'input_blocks.0.0.weight'}
|
| ]
|
| )
|
| for key, value in ic_model_state_dict.items()
|
| }
|
| )
|
|
|
| return (work_model, ic_model) |
