315 lines
11 KiB
Python
315 lines
11 KiB
Python
|
|
import torch
|
|
import numpy as np
|
|
from torchvision.transforms import ToTensor
|
|
from diffusers import AutoencoderKL
|
|
from diffusers import ControlNetModel
|
|
from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, StableDiffusionXLControlNetPipeline, StableDiffusionControlNetPipeline
|
|
from diffusers import StableDiffusionXLInpaintPipeline, StableDiffusionInpaintPipeline, StableDiffusionXLControlNetInpaintPipeline, StableDiffusionControlNetInpaintPipeline
|
|
from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionXLControlNetImg2ImgPipeline, StableDiffusionControlNetImg2ImgPipeline
|
|
from diffusers.utils import load_image
|
|
import folder_paths
|
|
from PIL import Image, ImageOps
|
|
from abc import ABC, abstractmethod
|
|
|
|
class ImageInference(ABC):
|
|
def setup_generator(self, seed):
|
|
return torch.Generator(device='cuda').manual_seed(seed)
|
|
|
|
@abstractmethod
|
|
def infer_image(self, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
class Text2ImageInference(ImageInference):
|
|
def infer_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, controlnet_image=None, controlnet_scale=None):
|
|
generator = self.setup_generator(seed)
|
|
|
|
args = {
|
|
"prompt": positive,
|
|
"generator": generator,
|
|
"num_inference_steps": steps,
|
|
"guidance_scale": cfg,
|
|
"width": width,
|
|
"height": height,
|
|
}
|
|
|
|
if negative != '':
|
|
args['negative_prompt'] = negative
|
|
|
|
if controlnet_image:
|
|
args['image'] = controlnet_image
|
|
args['controlnet_conditioning_scale'] = float(controlnet_scale)
|
|
|
|
|
|
images = pipeline(**args).images
|
|
|
|
return (convert_images_to_tensors(images))
|
|
|
|
class InpaintInference(ImageInference):
|
|
def infer_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, input_image, mask_image, mask_invert, controlnet_image=None, controlnet_scale=None, ):
|
|
generator = self.setup_generator(seed)
|
|
|
|
args = {
|
|
"prompt": positive,
|
|
"generator": generator,
|
|
"num_inference_steps": steps,
|
|
"guidance_scale": cfg,
|
|
"width": width,
|
|
"height": height,
|
|
|
|
}
|
|
|
|
if negative != '':
|
|
args['negative_prompt'] = negative
|
|
|
|
if controlnet_image:
|
|
args['control_image'] = controlnet_image
|
|
args['controlnet_conditioning_scale'] = float(controlnet_scale)
|
|
|
|
|
|
input_image = load_image(input_image)
|
|
args['image']= input_image
|
|
|
|
|
|
mask_image = load_image(mask_image)
|
|
if mask_invert:
|
|
mask_image = invert_mask(mask_image)
|
|
|
|
args['mask_image'] = mask_image
|
|
|
|
images = pipeline(**args).images
|
|
|
|
return (convert_images_to_tensors(images))
|
|
|
|
class Img2ImgInference(ImageInference):
|
|
def infer_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, input_image, controlnet_image=None, controlnet_scale=None ):
|
|
generator = self.setup_generator(seed)
|
|
|
|
args = {
|
|
"prompt": positive,
|
|
"generator": generator,
|
|
"num_inference_steps": steps,
|
|
"guidance_scale": cfg,
|
|
"width": width,
|
|
"height": height,
|
|
}
|
|
|
|
if negative != '':
|
|
args['negative_prompt'] = negative
|
|
|
|
if controlnet_image:
|
|
args['control_image'] = controlnet_image
|
|
args['controlnet_conditioning_scale'] = float(controlnet_scale)
|
|
|
|
if input_image != '':
|
|
input_image = load_image(input_image)
|
|
args['image']= input_image
|
|
|
|
images = pipeline(**args).images
|
|
|
|
return (convert_images_to_tensors(images))
|
|
|
|
|
|
|
|
class PipelineFactory:
|
|
"""
|
|
A factory class for creating pipelines based on the provided pipeline type.
|
|
"""
|
|
|
|
_pipeline_creators = {}
|
|
|
|
@classmethod
|
|
def register_creator(cls, pipeline_type, pipeline_creator_class):
|
|
"""
|
|
Register a pipeline creator class for a specific pipeline type.
|
|
|
|
Args:
|
|
pipeline_type (str): The type of pipeline to register the creator for.
|
|
pipeline_creator_class (class): The class responsible for creating the pipeline.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
cls._pipeline_creators[pipeline_type] = pipeline_creator_class
|
|
|
|
|
|
def get_pipeline(self, low_vram, pipeline_type, model, vae, controlnet_model, is_sdxl, torch_dtype ):
|
|
"""
|
|
Get a pipeline based on the provided parameters and set it to hardware.
|
|
Args:
|
|
low_vram (bool): Flag indicating low VRAM availability.
|
|
pipeline_type (str): The type of pipeline to create.
|
|
model: The model for the pipeline.
|
|
vae: The VAE for the pipeline.
|
|
controlnet_model: The controlnet model for the pipeline.
|
|
is_sdxl (bool): Flag indicating if SDXL is used.
|
|
torch_dtype: The torch data type.
|
|
|
|
Returns:
|
|
The created and set pipeline.
|
|
"""
|
|
# Check if pipeline is registered in the pipeline registry
|
|
if pipeline_type not in self._pipeline_creators:
|
|
raise ValueError(f"Unsupported pipeline type: {pipeline_type}")
|
|
|
|
# Get the right pipeline creator class
|
|
pipeline_creator_class = self._pipeline_creators[pipeline_type]
|
|
|
|
# Instantiate the right pipeline by sending the arguments to its constructor
|
|
pipeline_creator = pipeline_creator_class(model, vae, controlnet_model, is_sdxl, torch_dtype)
|
|
|
|
# Initialize the instantiated pipeline
|
|
pipeline = pipeline_creator.initialize_pipeline()
|
|
|
|
# Set the pipeline to hardware
|
|
return self.set_pipeline(low_vram, pipeline)
|
|
|
|
|
|
def set_pipeline(self, low_vram, pipeline):
|
|
"""
|
|
Set the pipeline to hardware and return it.
|
|
Args:
|
|
low_vram (bool): Flag indicating low VRAM availability.
|
|
pipeline: The pipeline to set to hardware.
|
|
|
|
Returns:
|
|
The pipeline set to hardware.
|
|
"""
|
|
if low_vram:
|
|
pipeline.enable_model_cpu_offload()
|
|
|
|
device = 'cpu' if low_vram or not torch.cuda.is_available() else 'cuda'
|
|
pipeline = pipeline.to(device)
|
|
|
|
print(f'pipeline is set to {device}')
|
|
|
|
return pipeline
|
|
|
|
class PipelineCreator(ABC):
|
|
"""
|
|
An abstract base class for creating pipeline creators with specific attributes.
|
|
"""
|
|
def __init__(self, model, vae, controlnet_model, is_sdxl, torch_dtype):
|
|
self.model = model
|
|
self.vae = vae
|
|
self.controlnet_model = controlnet_model
|
|
self.torch_dtype = torch_dtype
|
|
self.is_sdxl = is_sdxl
|
|
|
|
@abstractmethod
|
|
def initialize_pipeline(self):
|
|
raise NotImplementedError
|
|
|
|
class Text2ImgPipelineCreator(PipelineCreator):
|
|
"""
|
|
A PipelineCreator dedicated to initizalizing a Text2Img pipeline.
|
|
"""
|
|
|
|
def initialize_pipeline(self):
|
|
args = {
|
|
"pretrained_model_link_or_path": folder_paths.get_full_path("checkpoints", self.model),
|
|
"torch_dtype": self.torch_dtype
|
|
}
|
|
|
|
if self.vae != '':
|
|
args['vae'] = AutoencoderKL.from_pretrained(self.vae, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.controlnet_model != '':
|
|
args['controlnet'] = ControlNetModel.from_pretrained(self.controlnet_model, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.is_sdxl:
|
|
pipeline = StableDiffusionXLControlNetPipeline if self.controlnet_model != '' else StableDiffusionXLPipeline # Loads the ControlNet Pipe if controlnet is not empty
|
|
else:
|
|
pipeline = StableDiffusionControlNetPipeline if self.controlnet_model != '' else StableDiffusionPipeline
|
|
|
|
return pipeline.from_single_file(**args)
|
|
|
|
class Img2ImgPipelineCreator(PipelineCreator):
|
|
"""
|
|
A PipelineCreator dedicated to initizalizing a Img2Img pipeline.
|
|
"""
|
|
|
|
def initialize_pipeline(self):
|
|
args = {
|
|
"pretrained_model_link_or_path": folder_paths.get_full_path("checkpoints", self.model),
|
|
"torch_dtype": self.torch_dtype
|
|
}
|
|
|
|
if self.vae != '':
|
|
args['vae'] = AutoencoderKL.from_pretrained(self.vae, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.controlnet_model != '':
|
|
print('a controlnet was detected')
|
|
args['controlnet'] = ControlNetModel.from_pretrained(self.controlnet_model, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.is_sdxl:
|
|
pipeline = StableDiffusionXLControlNetImg2ImgPipeline if self.controlnet_model != '' else StableDiffusionXLImg2ImgPipeline
|
|
else:
|
|
pipeline = StableDiffusionControlNetImg2ImgPipeline if self.controlnet_model != '' else StableDiffusionImg2ImgPipeline
|
|
|
|
return pipeline.from_single_file(**args)
|
|
|
|
class InpaintPipelineCreator(PipelineCreator):
|
|
"""
|
|
A PipelineCreator dedicated to initizalizing an Inpaint pipeline.
|
|
"""
|
|
|
|
def initialize_pipeline(self):
|
|
args = {
|
|
"pretrained_model_link_or_path": folder_paths.get_full_path("checkpoints", self.model),
|
|
"torch_dtype": self.torch_dtype
|
|
}
|
|
|
|
if self.vae != '':
|
|
args['vae'] = AutoencoderKL.from_pretrained(self.vae, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.controlnet_model != '':
|
|
args['controlnet'] = ControlNetModel.from_pretrained(self.controlnet_model, torch_dtype=self.torch_dtype, use_safetensors=True)
|
|
|
|
if self.is_sdxl:
|
|
pipeline = StableDiffusionXLControlNetInpaintPipeline if self.controlnet_model != '' else StableDiffusionXLInpaintPipeline
|
|
else:
|
|
pipeline = StableDiffusionControlNetInpaintPipeline if self.controlnet_model != '' else StableDiffusionInpaintPipeline
|
|
|
|
return pipeline.from_single_file(**args)
|
|
|
|
|
|
|
|
|
|
def convert_images_to_tensors(images):
|
|
return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
|
|
|
|
def convert_tensors_to_images(images: torch.tensor):
|
|
return [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images]
|
|
|
|
|
|
def is_belong_to_blocks(key, blocks):
|
|
try:
|
|
for g in blocks:
|
|
if g in key:
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
raise type(e)(f'failed to is_belong_to_block, due to: {e}')
|
|
|
|
def filter_lora(state_dict, blocks_):
|
|
try:
|
|
return {k: v for k, v in state_dict.items() if is_belong_to_blocks(k, blocks_)}
|
|
except Exception as e:
|
|
raise type(e)(f'failed to filter_lora, due to: {e}')
|
|
|
|
def scale_lora(state_dict, alpha):
|
|
try:
|
|
return {k: v * alpha for k, v in state_dict.items()}
|
|
except Exception as e:
|
|
raise type(e)(f'failed to scale_lora, due to: {e}')
|
|
|
|
|
|
def invert_mask(image):
|
|
mask = image.convert('L')
|
|
|
|
return ImageOps.invert(mask)
|
|
|
|
|