Add files via upload
This commit is contained in:
GiusTex
+13
@@ -0,0 +1,13 @@
|
||||
from .nodes import (PadImageForDiffusersOutpaint, LoadDiffusersOutpaintModels, DiffusersImageOutpaint)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PadImageForDiffusersOutpaint": PadImageForDiffusersOutpaint,
|
||||
"LoadDiffusersOutpaintModels": LoadDiffusersOutpaintModels,
|
||||
"DiffusersImageOutpaint": DiffusersImageOutpaint
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PadImageForDiffusersOutpaint": "Pad Image For Diffusers Outpaint",
|
||||
"LoadDiffusersOutpaintModels": "Load Diffusers Outpaint Models",
|
||||
"DiffusersImageOutpaint": "Diffusers Image Outpaint"
|
||||
}
|
||||
+1085
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,296 @@
|
||||
import torch
|
||||
import gc
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from folder_paths import map_legacy, folder_names_and_paths
|
||||
from .controlnet_union import ControlNetModel_Union
|
||||
from .pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
|
||||
from diffusers import AutoencoderKL, TCDScheduler
|
||||
from diffusers.models.model_loading_utils import load_state_dict
|
||||
|
||||
|
||||
# Get the absolute path of various directories
|
||||
my_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
class PadImageForDiffusersOutpaint:
|
||||
_alignment_options = ["Middle", "Left", "Right", "Top", "Bottom"]
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"width": ("INT", {"default": 720, "min": 320, "max": 1536, "tooltip": "The width used for the image."}),
|
||||
"height": ("INT", {"default": 1280, "min": 320, "max": 1536, "tooltip": "The height used for the image."}),
|
||||
"alignment": (s._alignment_options, {"tooltip": "Where the original image should be in the outpainted one"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
||||
RETURN_NAMES = ("IMAGE", "MASK", "diffuser_outpaint_cnet_image")
|
||||
FUNCTION = "expand_image"
|
||||
CATEGORY = "DiffusersOutpaint"
|
||||
|
||||
def expand_image(self, image, width, height, alignment="Middle"):
|
||||
|
||||
# Resize Image
|
||||
def can_expand(source_width, source_height, target_width, target_height, alignment):
|
||||
"""Checks if the image can be expanded based on the alignment."""
|
||||
if alignment in ("Left", "Right") and source_width >= target_width:
|
||||
return False
|
||||
if alignment in ("Top", "Bottom") and source_height >= target_height:
|
||||
return False
|
||||
return True
|
||||
|
||||
im=tensor2pil(image)
|
||||
source=im.convert('RGB')
|
||||
target_size = (width, height)
|
||||
|
||||
# Upscale if source is smaller than target in both dimensions
|
||||
if source.width < target_size[0] and source.height < target_size[1]:
|
||||
scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
|
||||
new_width = int(source.width * scale_factor)
|
||||
new_height = int(source.height * scale_factor)
|
||||
source = source.resize((new_width, new_height), Image.LANCZOS)
|
||||
|
||||
if source.width > target_size[0] or source.height > target_size[1]:
|
||||
scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
|
||||
new_width = int(source.width * scale_factor)
|
||||
new_height = int(source.height * scale_factor)
|
||||
source = source.resize((new_width, new_height), Image.LANCZOS)
|
||||
|
||||
if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment):
|
||||
alignment = "Middle"
|
||||
# Calculate margins based on alignment
|
||||
if alignment == "Middle":
|
||||
margin_x = (target_size[0] - source.width) // 2
|
||||
margin_y = (target_size[1] - source.height) // 2
|
||||
elif alignment == "Left":
|
||||
margin_x = 0
|
||||
margin_y = (target_size[1] - source.height) // 2
|
||||
elif alignment == "Right":
|
||||
margin_x = target_size[0] - source.width
|
||||
margin_y = (target_size[1] - source.height) // 2
|
||||
elif alignment == "Top":
|
||||
margin_x = (target_size[0] - source.width) // 2
|
||||
margin_y = 0
|
||||
elif alignment == "Bottom":
|
||||
margin_x = (target_size[0] - source.width) // 2
|
||||
margin_y = target_size[1] - source.height
|
||||
|
||||
background = Image.new('RGB', target_size, (255, 255, 255))
|
||||
background.paste(source, (margin_x, margin_y))
|
||||
|
||||
image=pil2tensor(background)
|
||||
#----------------------------------------------------
|
||||
d1, d2, d3, d4 = image.size()
|
||||
left, top, bottom, right = 0, 0, 0, 0
|
||||
# Image
|
||||
new_image = torch.ones(
|
||||
(d1, d2 + top + bottom, d3 + left + right, d4),
|
||||
dtype=torch.float32,
|
||||
) * 0.5
|
||||
new_image[:, top:top + d2, left:left + d3, :] = image
|
||||
#----------------------------------------------------
|
||||
# Mask coordinates
|
||||
if alignment == "Middle":
|
||||
margin_x = (width - new_width) // 2
|
||||
margin_y = (height - new_height) // 2
|
||||
elif alignment == "Left":
|
||||
margin_x = 0
|
||||
margin_y = (height - new_height) // 2
|
||||
elif alignment == "Right":
|
||||
margin_x = width - new_width
|
||||
margin_y = (height - new_height) // 2
|
||||
elif alignment == "Top":
|
||||
margin_x = (width - new_width) // 2
|
||||
margin_y = 0
|
||||
elif alignment == "Bottom":
|
||||
margin_x = (width - new_width) // 2
|
||||
margin_y = height - new_height
|
||||
|
||||
# Create mask as big as new img
|
||||
mask = torch.ones(
|
||||
(height, width),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
# Create hole in mask
|
||||
t = torch.zeros(
|
||||
(new_height, new_width),
|
||||
dtype=torch.float32
|
||||
)
|
||||
# Create holed mask
|
||||
mask[margin_y:margin_y + new_height,
|
||||
margin_x:margin_x + new_width
|
||||
] = t
|
||||
#----------------------------------------------------
|
||||
# Prepare "cn_image" for diffusers outpaint
|
||||
im=tensor2pil(new_image)
|
||||
pil_new_image=im.convert('RGB')
|
||||
|
||||
pil_mask=tensor2pil(mask)
|
||||
|
||||
cnet_image = pil_new_image.copy() # copy background as cnet_image
|
||||
cnet_image.paste(0, (0, 0), pil_mask) # paste mask over cnet_image, cropping it a bit
|
||||
|
||||
tensor_cnet_image=pil2tensor(cnet_image)
|
||||
|
||||
return (new_image, mask, tensor_cnet_image,)
|
||||
|
||||
|
||||
def get_first_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]:
|
||||
folder_name = map_legacy(folder_name)
|
||||
global folder_names_and_paths
|
||||
folders = folder_names_and_paths[folder_name]
|
||||
root_folder = folders[0][0]
|
||||
visible_folders = [name for name in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, name))]
|
||||
return visible_folders
|
||||
|
||||
class LoadDiffusersOutpaintModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (get_first_folder_list("unet"), {"tooltip": "The diffuser model used for denoising the input latent. (Put model files in the unet folder)."}),
|
||||
"vae": (get_first_folder_list("vae"), {"tooltip": "The vae model used for denoising the input latent.(Put model files in the vae folder)."}),
|
||||
"controlnet_model": (get_first_folder_list("controlnet"), {"tooltip": "The controlnet model used for denoising the input latent.(Put model files in the controlnet folder)."}),
|
||||
},
|
||||
"optional": {
|
||||
"keep_models_in_vram": ("BOOLEAN", {"default": False, "tooltip": "Set to false to unload diffusion models, and maybe others too, from vram."}),
|
||||
"enable_model_cpu_offload": ("BOOLEAN", {"default": True, "tooltip": "Reduces memory usage with a low impact on performance."}),
|
||||
"enable_vae_slicing": ("BOOLEAN", {"default": True, "tooltip": "VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes."}),
|
||||
"enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PIPE",)
|
||||
RETURN_NAMES = ("diffusers_outpaint_pipe",)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "DiffusersOutpaint"
|
||||
|
||||
def load(self, model, vae, controlnet_model, keep_models_in_vram, enable_model_cpu_offload, enable_vae_slicing, enable_vae_tiling):
|
||||
# Go 2 folders back
|
||||
comfy_dir = os.path.dirname(os.path.dirname(my_dir))
|
||||
|
||||
model_path = f"{comfy_dir}/models/unet/{model}"
|
||||
vae_path = f"{comfy_dir}/models/vae/{vae}"
|
||||
controlnet_path = f"{comfy_dir}/models/controlnet/{controlnet_model}"
|
||||
#-----------------------------------------------------------------------
|
||||
|
||||
# Set up Controlnet-Union-Promax-SDXL model
|
||||
config_file = f"{controlnet_path}/config_promax.json"
|
||||
config = ControlNetModel_Union.load_config(config_file)
|
||||
controlnet_model = ControlNetModel_Union.from_config(config)
|
||||
|
||||
model_file = f"{controlnet_path}/diffusion_pytorch_model_promax.safetensors"
|
||||
state_dict = load_state_dict(model_file)
|
||||
|
||||
model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
|
||||
controlnet_model, state_dict, model_file, f"{controlnet_path}"
|
||||
)
|
||||
model.to(device="cuda", dtype=torch.float16)
|
||||
#-----------------------------------------------------------------------
|
||||
|
||||
# Set up VAE
|
||||
vae = AutoencoderKL.from_pretrained(f"{vae_path}", torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
if enable_vae_slicing:
|
||||
vae.enable_slicing()
|
||||
else:
|
||||
vae.disable_slicing()
|
||||
|
||||
if enable_vae_tiling:
|
||||
vae.enable_tiling()
|
||||
else:
|
||||
vae.disable_tiling()
|
||||
#-----------------------------------------------------------------------
|
||||
|
||||
# Load Controlnet + Vae into RealVisXL model
|
||||
pipe = StableDiffusionXLFillPipeline.from_pretrained(
|
||||
f"{model_path}",
|
||||
torch_dtype=torch.float16,
|
||||
vae=vae,
|
||||
controlnet=controlnet_model,
|
||||
variant="fp16",
|
||||
).to("cuda")
|
||||
|
||||
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
diffusers_outpaint_pipe = {
|
||||
"pipe": pipe,
|
||||
"vae": vae,
|
||||
"model": model,
|
||||
"controlnet_model": controlnet_model,
|
||||
"state_dict": state_dict,
|
||||
"model_file": model_file,
|
||||
"enable_model_cpu_offload": enable_model_cpu_offload,
|
||||
"keep_models_in_vram": keep_models_in_vram
|
||||
}
|
||||
|
||||
return (diffusers_outpaint_pipe,)
|
||||
|
||||
|
||||
# Tensor to PIL (grabbed from WAS Suite)
|
||||
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
||||
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
||||
|
||||
# Convert PIL to Tensor (grabbed from WAS Suite)
|
||||
def pil2tensor(image: Image.Image) -> torch.Tensor:
|
||||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
class DiffusersImageOutpaint:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"diffusers_outpaint_pipe": ("PIPE", {"tooltip": "Load the diffusers outpaint models."}),
|
||||
"diffuser_outpaint_cnet_image": ("IMAGE", {"tooltip": "The image to outpaint."}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Fake seed, workaround used to keep generating different outpaints. Set to -1 to generate different images, or a fixed number to stop that."}),
|
||||
"steps": ("INT", {"default": 8, "min": 4, "max": 20, "tooltip": "The number of steps used in the denoising process."}),
|
||||
"extra_prompt": ("STRING", {"default": "", "tooltip": "The extra prompt to append, describing attributes etc. you want to include in the image. Default: \"(extra_prompt), high quality, 4k\""}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "DiffusersOutpaint"
|
||||
|
||||
def sample(self, diffusers_outpaint_pipe, diffuser_outpaint_cnet_image, seed, steps, extra_prompt=None):
|
||||
pipe = diffusers_outpaint_pipe["pipe"]
|
||||
|
||||
final_prompt = f"{extra_prompt}, high quality, 4k"
|
||||
|
||||
cnet_image = diffuser_outpaint_cnet_image
|
||||
cnet_image=tensor2pil(cnet_image)
|
||||
cnet_image=cnet_image.convert('RGB')
|
||||
|
||||
(prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
) = pipe.encode_prompt(final_prompt, "cuda", True)
|
||||
|
||||
if diffusers_outpaint_pipe["enable_model_cpu_offload"]:
|
||||
pipe.enable_model_cpu_offload()
|
||||
|
||||
generated_images = list(pipe(
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
image=cnet_image,
|
||||
num_inference_steps=steps
|
||||
))
|
||||
|
||||
if not diffusers_outpaint_pipe["keep_models_in_vram"]:
|
||||
del pipe, diffusers_outpaint_pipe["vae"], diffusers_outpaint_pipe["model"], diffusers_outpaint_pipe["controlnet_model"], diffusers_outpaint_pipe["state_dict"], diffusers_outpaint_pipe["model_file"], prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
last_image = generated_images[-1] # Access the last image
|
||||
image = last_image.convert("RGBA")
|
||||
output=pil2tensor(image)
|
||||
|
||||
return (output,)
|
||||
@@ -0,0 +1,559 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import cv2
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
|
||||
|
||||
from .controlnet_union import ControlNetModel_Union
|
||||
|
||||
|
||||
def latents_to_rgb(latents):
|
||||
weights = ((60, -60, 25, -70), (60, -5, 15, -50), (60, 10, -5, -35))
|
||||
|
||||
weights_tensor = torch.t(
|
||||
torch.tensor(weights, dtype=latents.dtype).to(latents.device)
|
||||
)
|
||||
biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(
|
||||
latents.device
|
||||
)
|
||||
rgb_tensor = torch.einsum(
|
||||
"...lxy,lr -> ...rxy", latents, weights_tensor
|
||||
) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
|
||||
image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy()
|
||||
image_array = image_array.transpose(1, 2, 0) # Change the order of dimensions
|
||||
|
||||
denoised_image = cv2.fastNlMeansDenoisingColored(image_array, None, 10, 10, 7, 21)
|
||||
blurred_image = cv2.GaussianBlur(denoised_image, (5, 5), 0)
|
||||
final_image = PIL.Image.fromarray(blurred_image)
|
||||
|
||||
width, height = final_image.size
|
||||
final_image = final_image.resize(
|
||||
(width * 8, height * 8), PIL.Image.Resampling.LANCZOS
|
||||
)
|
||||
|
||||
return final_image
|
||||
|
||||
|
||||
def retrieve_timesteps(
|
||||
scheduler,
|
||||
num_inference_steps: Optional[int] = None,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
|
||||
return timesteps, num_inference_steps
|
||||
|
||||
|
||||
class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin):
|
||||
model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
|
||||
_optional_components = [
|
||||
"tokenizer",
|
||||
"tokenizer_2",
|
||||
"text_encoder",
|
||||
"text_encoder_2",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
text_encoder_2: CLIPTextModelWithProjection,
|
||||
tokenizer: CLIPTokenizer,
|
||||
tokenizer_2: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
controlnet: ControlNetModel_Union,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
force_zeros_for_empty_prompt: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
text_encoder_2=text_encoder_2,
|
||||
tokenizer=tokenizer,
|
||||
tokenizer_2=tokenizer_2,
|
||||
unet=unet,
|
||||
controlnet=controlnet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
|
||||
)
|
||||
self.control_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor,
|
||||
do_convert_rgb=True,
|
||||
do_normalize=False,
|
||||
)
|
||||
|
||||
self.register_to_config(
|
||||
force_zeros_for_empty_prompt=force_zeros_for_empty_prompt
|
||||
)
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str,
|
||||
device: Optional[torch.device] = None,
|
||||
do_classifier_free_guidance: bool = True,
|
||||
):
|
||||
device = device or self._execution_device
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
if prompt is not None:
|
||||
batch_size = len(prompt)
|
||||
|
||||
# Define tokenizers and text encoders
|
||||
tokenizers = (
|
||||
[self.tokenizer, self.tokenizer_2]
|
||||
if self.tokenizer is not None
|
||||
else [self.tokenizer_2]
|
||||
)
|
||||
text_encoders = (
|
||||
[self.text_encoder, self.text_encoder_2]
|
||||
if self.text_encoder is not None
|
||||
else [self.text_encoder_2]
|
||||
)
|
||||
|
||||
prompt_2 = prompt
|
||||
prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
|
||||
|
||||
# textual inversion: process multi-vector tokens if necessary
|
||||
prompt_embeds_list = []
|
||||
prompts = [prompt, prompt_2]
|
||||
for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
|
||||
text_inputs = tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=tokenizer.model_max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
text_input_ids = text_inputs.input_ids
|
||||
|
||||
prompt_embeds = text_encoder(
|
||||
text_input_ids.to(device), output_hidden_states=True
|
||||
)
|
||||
|
||||
# We are only ALWAYS interested in the pooled output of the final text encoder
|
||||
pooled_prompt_embeds = prompt_embeds[0]
|
||||
prompt_embeds = prompt_embeds.hidden_states[-2]
|
||||
prompt_embeds_list.append(prompt_embeds)
|
||||
|
||||
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
|
||||
|
||||
# get unconditional embeddings for classifier free guidance
|
||||
zero_out_negative_prompt = True
|
||||
negative_prompt_embeds = None
|
||||
negative_pooled_prompt_embeds = None
|
||||
|
||||
if do_classifier_free_guidance and zero_out_negative_prompt:
|
||||
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
|
||||
negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
|
||||
elif do_classifier_free_guidance and negative_prompt_embeds is None:
|
||||
negative_prompt = ""
|
||||
negative_prompt_2 = negative_prompt
|
||||
|
||||
# normalize str to list
|
||||
negative_prompt = (
|
||||
batch_size * [negative_prompt]
|
||||
if isinstance(negative_prompt, str)
|
||||
else negative_prompt
|
||||
)
|
||||
negative_prompt_2 = (
|
||||
batch_size * [negative_prompt_2]
|
||||
if isinstance(negative_prompt_2, str)
|
||||
else negative_prompt_2
|
||||
)
|
||||
|
||||
uncond_tokens: List[str]
|
||||
if prompt is not None and type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
else:
|
||||
uncond_tokens = [negative_prompt, negative_prompt_2]
|
||||
|
||||
negative_prompt_embeds_list = []
|
||||
for negative_prompt, tokenizer, text_encoder in zip(
|
||||
uncond_tokens, tokenizers, text_encoders
|
||||
):
|
||||
max_length = prompt_embeds.shape[1]
|
||||
uncond_input = tokenizer(
|
||||
negative_prompt,
|
||||
padding="max_length",
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
negative_prompt_embeds = text_encoder(
|
||||
uncond_input.input_ids.to(device),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
# We are only ALWAYS interested in the pooled output of the final text encoder
|
||||
negative_pooled_prompt_embeds = negative_prompt_embeds[0]
|
||||
negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
|
||||
|
||||
negative_prompt_embeds_list.append(negative_prompt_embeds)
|
||||
|
||||
negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
|
||||
|
||||
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
|
||||
|
||||
bs_embed, seq_len, _ = prompt_embeds.shape
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
prompt_embeds = prompt_embeds.repeat(1, 1, 1)
|
||||
prompt_embeds = prompt_embeds.view(bs_embed * 1, seq_len, -1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
||||
seq_len = negative_prompt_embeds.shape[1]
|
||||
|
||||
if self.text_encoder_2 is not None:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(
|
||||
dtype=self.text_encoder_2.dtype, device=device
|
||||
)
|
||||
else:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(
|
||||
dtype=self.unet.dtype, device=device
|
||||
)
|
||||
|
||||
negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1)
|
||||
negative_prompt_embeds = negative_prompt_embeds.view(
|
||||
batch_size * 1, seq_len, -1
|
||||
)
|
||||
|
||||
pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
|
||||
if do_classifier_free_guidance:
|
||||
negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(
|
||||
1, 1
|
||||
).view(bs_embed * 1, -1)
|
||||
|
||||
return (
|
||||
prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
)
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
image,
|
||||
controlnet_conditioning_scale=1.0,
|
||||
):
|
||||
if prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide `prompt_embeds`. Cannot leave `prompt_embeds` undefined."
|
||||
)
|
||||
|
||||
if negative_prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide `negative_prompt_embeds`. Cannot leave `negative_prompt_embeds` undefined."
|
||||
)
|
||||
|
||||
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
||||
raise ValueError(
|
||||
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
||||
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
||||
f" {negative_prompt_embeds.shape}."
|
||||
)
|
||||
|
||||
if prompt_embeds is not None and pooled_prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
|
||||
)
|
||||
|
||||
if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
|
||||
)
|
||||
|
||||
# Check `image`
|
||||
is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
|
||||
self.controlnet, torch._dynamo.eval_frame.OptimizedModule
|
||||
)
|
||||
if (
|
||||
isinstance(self.controlnet, ControlNetModel_Union)
|
||||
or is_compiled
|
||||
and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
|
||||
):
|
||||
if not isinstance(image, PIL.Image.Image):
|
||||
raise TypeError(
|
||||
f"image must be passed and has to be a PIL image, but is {type(image)}"
|
||||
)
|
||||
|
||||
else:
|
||||
assert False
|
||||
|
||||
# Check `controlnet_conditioning_scale`
|
||||
if (
|
||||
isinstance(self.controlnet, ControlNetModel_Union)
|
||||
or is_compiled
|
||||
and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
|
||||
):
|
||||
if not isinstance(controlnet_conditioning_scale, float):
|
||||
raise TypeError(
|
||||
"For single controlnet: `controlnet_conditioning_scale` must be type `float`."
|
||||
)
|
||||
else:
|
||||
assert False
|
||||
|
||||
def prepare_image(self, image, device, dtype, do_classifier_free_guidance=False):
|
||||
image = self.control_image_processor.preprocess(image).to(dtype=torch.float32)
|
||||
|
||||
image_batch_size = image.shape[0]
|
||||
|
||||
image = image.repeat_interleave(image_batch_size, dim=0)
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
image = torch.cat([image] * 2)
|
||||
|
||||
return image
|
||||
|
||||
def prepare_latents(
|
||||
self, batch_size, num_channels_latents, height, width, dtype, device
|
||||
):
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
int(height) // self.vae_scale_factor,
|
||||
int(width) // self.vae_scale_factor,
|
||||
)
|
||||
|
||||
latents = randn_tensor(shape, device=device, dtype=dtype)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
@property
|
||||
def do_classifier_free_guidance(self):
|
||||
return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt_embeds: torch.Tensor,
|
||||
negative_prompt_embeds: torch.Tensor,
|
||||
pooled_prompt_embeds: torch.Tensor,
|
||||
negative_pooled_prompt_embeds: torch.Tensor,
|
||||
image: PipelineImageInput = None,
|
||||
num_inference_steps: int = 8,
|
||||
guidance_scale: float = 1.5,
|
||||
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
|
||||
):
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(
|
||||
prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
image,
|
||||
controlnet_conditioning_scale,
|
||||
)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
|
||||
# 2. Define call parameters
|
||||
batch_size = 1
|
||||
device = self._execution_device
|
||||
|
||||
# 4. Prepare image
|
||||
if isinstance(self.controlnet, ControlNetModel_Union):
|
||||
image = self.prepare_image(
|
||||
image=image,
|
||||
device=device,
|
||||
dtype=self.controlnet.dtype,
|
||||
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
||||
)
|
||||
height, width = image.shape[-2:]
|
||||
else:
|
||||
assert False
|
||||
|
||||
# 5. Prepare timesteps
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler, num_inference_steps, device
|
||||
)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# 6. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
)
|
||||
|
||||
# 7 Prepare added time ids & embeddings
|
||||
add_text_embeds = pooled_prompt_embeds
|
||||
|
||||
add_time_ids = negative_add_time_ids = torch.tensor(
|
||||
image.shape[-2:] + torch.Size([0, 0]) + image.shape[-2:]
|
||||
).unsqueeze(0)
|
||||
|
||||
if self.do_classifier_free_guidance:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
add_text_embeds = torch.cat(
|
||||
[negative_pooled_prompt_embeds, add_text_embeds], dim=0
|
||||
)
|
||||
add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
|
||||
|
||||
prompt_embeds = prompt_embeds.to(device)
|
||||
add_text_embeds = add_text_embeds.to(device)
|
||||
add_time_ids = add_time_ids.to(device).repeat(batch_size, 1)
|
||||
|
||||
controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
|
||||
union_control_type = (
|
||||
torch.Tensor([0, 0, 0, 0, 0, 0, 1, 0])
|
||||
.to(device, dtype=prompt_embeds.dtype)
|
||||
.repeat(batch_size * 2, 1)
|
||||
)
|
||||
|
||||
added_cond_kwargs = {
|
||||
"text_embeds": add_text_embeds,
|
||||
"time_ids": add_time_ids,
|
||||
"control_type": union_control_type,
|
||||
}
|
||||
|
||||
controlnet_prompt_embeds = prompt_embeds
|
||||
controlnet_added_cond_kwargs = added_cond_kwargs
|
||||
|
||||
# 8. Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = (
|
||||
torch.cat([latents] * 2)
|
||||
if self.do_classifier_free_guidance
|
||||
else latents
|
||||
)
|
||||
latent_model_input = self.scheduler.scale_model_input(
|
||||
latent_model_input, t
|
||||
)
|
||||
|
||||
# controlnet(s) inference
|
||||
control_model_input = latent_model_input
|
||||
|
||||
down_block_res_samples, mid_block_res_sample = self.controlnet(
|
||||
control_model_input,
|
||||
t,
|
||||
encoder_hidden_states=controlnet_prompt_embeds,
|
||||
controlnet_cond_list=controlnet_image_list,
|
||||
conditioning_scale=controlnet_conditioning_scale,
|
||||
guess_mode=False,
|
||||
added_cond_kwargs=controlnet_added_cond_kwargs,
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
timestep_cond=None,
|
||||
cross_attention_kwargs={},
|
||||
down_block_additional_residuals=down_block_res_samples,
|
||||
mid_block_additional_residual=mid_block_res_sample,
|
||||
added_cond_kwargs=added_cond_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# perform guidance
|
||||
if self.do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(
|
||||
noise_pred, t, latents, return_dict=False
|
||||
)[0]
|
||||
|
||||
if i == 2:
|
||||
prompt_embeds = prompt_embeds[-1:]
|
||||
add_text_embeds = add_text_embeds[-1:]
|
||||
add_time_ids = add_time_ids[-1:]
|
||||
union_control_type = union_control_type[-1:]
|
||||
|
||||
added_cond_kwargs = {
|
||||
"text_embeds": add_text_embeds,
|
||||
"time_ids": add_time_ids,
|
||||
"control_type": union_control_type,
|
||||
}
|
||||
|
||||
controlnet_prompt_embeds = prompt_embeds
|
||||
controlnet_added_cond_kwargs = added_cond_kwargs
|
||||
|
||||
image = image[-1:]
|
||||
controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
|
||||
|
||||
self._guidance_scale = 0.0
|
||||
|
||||
if i == len(timesteps) - 1 or (
|
||||
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
yield latents_to_rgb(latents)
|
||||
|
||||
latents = latents / self.vae.config.scaling_factor
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = self.image_processor.postprocess(image)[0]
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
yield image
|
||||
@@ -0,0 +1,7 @@
|
||||
torch
|
||||
numpy==1.26.4
|
||||
transformers
|
||||
accelerate
|
||||
diffusers
|
||||
fastapi<0.113.0
|
||||
opencv-python
|
||||
Reference in New Issue
Block a user