272 lines
9.8 KiB
Python
272 lines
9.8 KiB
Python
import abc
|
|
import torch
|
|
import numpy as np
|
|
import comfy.model_management as mm
|
|
|
|
from PIL import Image
|
|
from diffusers import StableDiffusionPipeline
|
|
from tqdm import tqdm
|
|
|
|
from .utils import register_attention_control, AttentionStore, aggregate_attention, \
|
|
text_under_image, convert_preview_image, LOW_RESOURCE
|
|
|
|
|
|
class HFModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (
|
|
["runwayml/stable-diffusion-v1-5"],
|
|
{"default": "runwayml/stable-diffusion-v1-5"},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "TOKENIZER",)
|
|
RETURN_NAMES = ("model", "tokenizer",)
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def load_checkpoint(self, model):
|
|
device = mm.get_torch_device()
|
|
model = StableDiffusionPipeline.from_pretrained(model).to(device)
|
|
tokenizer = model.tokenizer
|
|
return (model, tokenizer)
|
|
|
|
|
|
class ShowImages:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":
|
|
{"image": ("IMAGE",)}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "view_images"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def view_images(self, image, num_rows=1, offset_ratio=0.02):
|
|
images = image
|
|
if type(images) is list:
|
|
num_empty = len(images) % num_rows
|
|
elif images.ndim == 4:
|
|
num_empty = images.shape[0] % num_rows
|
|
else:
|
|
images = [images]
|
|
num_empty = 0
|
|
|
|
empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
|
|
images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
|
|
num_items = len(images)
|
|
|
|
h, w, c = images[0].shape
|
|
offset = int(h * offset_ratio)
|
|
num_cols = num_items // num_rows
|
|
image_ = np.ones((h * num_rows + offset * (num_rows - 1),
|
|
w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
|
|
for i in range(num_rows):
|
|
for j in range(num_cols):
|
|
image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
|
|
i * num_cols + j]
|
|
pil_img = [Image.fromarray(image_)]
|
|
output_images = convert_preview_image(pil_img)
|
|
return (output_images,)
|
|
|
|
|
|
class Text2ImageInference:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"prompt": ("STRING", {"multiline": True, "default": "A painting of a squirrel eating a burger", }),
|
|
"steps": ("INT", {"default": 50, "min": 1, "max": 100, "step": 1, "display": "slider"}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 10.0, "step": 0.5, "display": "slider"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"latents": ("LATENTS",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE", "CONTROLLER", "LATENTS", "STRING")
|
|
RETURN_NAMES = ("vae", "controller", "latents", "prompt")
|
|
FUNCTION = "inference"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def init_latent(self, latent, model, height, width, generator, batch_size):
|
|
if latent is None:
|
|
latent = torch.randn(
|
|
(1, model.unet.in_channels, height // 8, width // 8),
|
|
generator=generator,
|
|
)
|
|
latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device)
|
|
return latent, latents
|
|
|
|
def diffusion_step(self, model, controller, latents, context, t, guidance_scale, low_resource=False):
|
|
if low_resource:
|
|
noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"]
|
|
noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"]
|
|
else:
|
|
latents_input = torch.cat([latents] * 2) # 40,4,64,64
|
|
noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)[
|
|
"sample"] # latents_input:80,4,64,64;41,77,768
|
|
noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
|
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
|
|
latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"]
|
|
latents = controller.step_callback(latents)
|
|
return latents
|
|
|
|
def inference(self, model, prompt, steps, guidance_scale, seed, latent=None):
|
|
# import pdb;
|
|
# pdb.set_trace()
|
|
low_resource = LOW_RESOURCE
|
|
prompt = [prompt]
|
|
|
|
controller = AttentionStore()
|
|
register_attention_control(model, controller)
|
|
|
|
height = width = 512
|
|
batch_size = len(prompt)
|
|
|
|
text_input = model.tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=model.tokenizer.model_max_length,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
)
|
|
text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]
|
|
max_length = text_input.input_ids.shape[-1]
|
|
uncond_input = model.tokenizer(
|
|
[""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
|
|
)
|
|
uncond_embeddings = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
|
|
|
|
context = [uncond_embeddings, text_embeddings]
|
|
if not low_resource:
|
|
context = torch.cat(context)
|
|
generator = torch.Generator().manual_seed(seed)
|
|
latent, latents = self.init_latent(latent, model, height, width, generator, batch_size)
|
|
|
|
# set timesteps
|
|
model.scheduler.set_timesteps(steps)
|
|
for t in tqdm(model.scheduler.timesteps):
|
|
latents = self.diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource)
|
|
|
|
return (model.vae, controller, latents, prompt)
|
|
|
|
|
|
class DecodeLatent:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"vae": ("VAE",),
|
|
"latents": ("LATENTS",)}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "decode"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def decode(self, vae, latents):
|
|
latents = 1 / 0.18215 * latents
|
|
image = vae.decode(latents)['sample']
|
|
image = (image / 2 + 0.5).clamp(0, 1)
|
|
image = image.cpu().permute(0, 2, 3, 1).numpy()
|
|
images = (image * 255).astype(np.uint8)
|
|
pil_images = [Image.fromarray(image) for image in images]
|
|
output_images = convert_preview_image(pil_images)
|
|
return (output_images,)
|
|
|
|
|
|
class ShowCrossAttn:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"controller": ("CONTROLLER",),
|
|
"prompt": ("STRING", {"forceInput": True, "default": ""}),
|
|
"tokenizer": ("TOKENIZER",),
|
|
"res": ("INT", {"default": 16, "min": 0, "max": 50, "step": 1, "display": "slider"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "show_crossattn"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def show_crossattn(self, controller, prompt, tokenizer, res, from_where=("up", "down"), select: int = 0):
|
|
attention_store = controller
|
|
tokens = tokenizer.encode(prompt[select])
|
|
decoder = tokenizer.decode
|
|
attention_maps = aggregate_attention(attention_store, prompt, res, from_where, True, select)
|
|
|
|
images = []
|
|
for i in range(len(tokens)):
|
|
image = attention_maps[:, :, i]
|
|
image = 255 * image / image.max()
|
|
image = image.unsqueeze(-1).expand(*image.shape, 3)
|
|
image = image.numpy().astype(np.uint8)
|
|
image = np.array(Image.fromarray(image).resize((256, 256)))
|
|
image = text_under_image(image, decoder(int(tokens[i])))
|
|
images.append(image)
|
|
images = np.stack(images, axis=0)
|
|
return (images,)
|
|
|
|
|
|
class ShowSelfAttn:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"controller": ("CONTROLLER",),
|
|
"prompt": ("STRING", {"forceInput": True, "default": ""}),
|
|
"tokenizer": ("TOKENIZER",),
|
|
"res": ("INT", {"default": 16, "min": 0, "max": 50, "step": 1, "display": "slider"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "show_selfattn"
|
|
CATEGORY = "VisualAttn"
|
|
|
|
def show_selfattn(self, controller, prompt, tokenizer, res, from_where=("up", "down"), select: int = 0, max_com=10):
|
|
attention_store = controller
|
|
tokens = tokenizer.encode(prompt[select])
|
|
decoder = tokenizer.decode
|
|
attention_maps = aggregate_attention(attention_store, prompt, res, from_where, False, select).numpy().reshape(
|
|
(res ** 2, res ** 2))
|
|
u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
|
|
|
|
images = []
|
|
for i in range(len(tokens)):
|
|
image = vh[i].reshape(res, res)
|
|
image = image - image.min()
|
|
image = 255 * image / image.max()
|
|
image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
|
|
image = np.array(Image.fromarray(image).resize((256, 256)))
|
|
image = text_under_image(image, decoder(int(tokens[i])))
|
|
images.append(image)
|
|
images = np.concatenate(images, axis=1)
|
|
return (images,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"HFModelLoader": HFModelLoader,
|
|
"ShowImages": ShowImages,
|
|
"Text2ImageInference": Text2ImageInference,
|
|
"DecodeLatent": DecodeLatent,
|
|
"ShowCrossAttn": ShowCrossAttn,
|
|
"ShowSelfAttn": ShowSelfAttn
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"HFModelLoader": "HF ModelLoader",
|
|
"ShowImages": "Show Images",
|
|
"Text2ImageInference": "Text2Image Inference",
|
|
"DecodeLatent": "Decode Latent",
|
|
"ShowCrossAttn": "Show CrossAttn Map",
|
|
"ShowSelfAttn": "Show SelfAttn Map"
|
|
}
|