Merge pull request #1831 from vladmandic/improve_diffusers_kandinsky_2

[Diffusers] Make all Kandinsky work
This commit is contained in:
Vladimir Mandic
2023-07-26 16:04:08 -04:00
committed by GitHub
4 changed files with 29 additions and 102 deletions
+2 -2
View File
@@ -20,7 +20,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
"force_download": False,
"resume_download": True,
"cache_dir": shared.opts.diffusers_dir,
# "use_auth_token": True,
"load_connected_pipeline": True,
}
if cache_dir is not None:
download_config["cache_dir"] = cache_dir
@@ -41,7 +41,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
model_info_dict = None
# some checkpoints need to be downloaded as "hidden" as they just serve as pre- or post-pipelines of other pipelines
if model_info_dict is not None and "prior" in model_info_dict:
download_dir = DiffusionPipeline.download(model_info_dict["prior"], **download_config)
download_dir = DiffusionPipeline.download(model_info_dict["prior"][0], **download_config)
model_info_dict["prior"] = download_dir
# mark prior as hidden
with open(os.path.join(download_dir, "hidden"), "w", encoding="utf-8") as f:
+9 -2
View File
@@ -9,6 +9,12 @@ from modules.lora_diffusers import lora_state, unload_diffusers_lora
from modules.processing import StableDiffusionProcessing
try:
import diffusers
except Exception as ex:
shared.log.error(f'Failed to import diffusers: {ex}')
def encode_prompt(encoder, prompt):
cfg = encoder.config
# TODO implement similar hijack for diffusers text encoder but following diffusers pipeline.encode_prompt concepts
@@ -40,7 +46,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
def set_pipeline_args(model, prompt, negative_prompt, **kwargs):
args = {}
pipeline = model.main if model.__class__.__name__ == 'PriorPipeline' else model
pipeline = model
signature = inspect.signature(type(pipeline).__call__)
possible = signature.parameters.keys()
generator_device = 'cpu' if shared.opts.diffusers_generator_device == "cpu" else shared.device
@@ -93,7 +99,8 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
shared.log.debug(f'Diffuser pipeline: {pipeline.__class__.__name__} task={sd_models.get_diffusers_task(model)} set={clean}')
return args
if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.sampler_name) and (p.sampler_name != 'Default'):
is_karras_compatible = shared.sd_model.__class__.__init__.__annotations__.get("scheduler", None) == diffusers.schedulers.scheduling_utils.KarrasDiffusionSchedulers
if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.sampler_name) and (p.sampler_name != 'Default') and is_karras_compatible:
sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None)
if sampler is None:
sampler = sd_samplers.all_samplers_map.get("UniPC")
+17 -97
View File
@@ -520,53 +520,6 @@ class ModelData:
model_data = ModelData()
class PriorPipeline:
def __init__(self, prior, main):
self.prior = prior
self.main = main
self.main.safety_checker = None
self.scheduler = main.scheduler
self.tokenizer = self.prior.tokenizer
self.unet = self.main.unet
def to(self, *args, **kwargs):
# only the prior is moved to CUDA in a first step
self.prior.to(*args, **kwargs)
def enable_model_cpu_offload(self, *args, **kwargs):
if hasattr(self.prior, 'enable_model_cpu_offload'):
self.prior.enable_model_cpu_offload(*args, **kwargs)
self.main.enable_model_cpu_offload(*args, **kwargs)
def enable_sequential_cpu_offload(self, *args, **kwargs):
if hasattr(self.prior, 'enable_sequential_cpu_offload'):
self.prior.enable_sequential_cpu_offload(*args, **kwargs)
self.main.enable_sequential_cpu_offload(*args, **kwargs)
def enable_xformers_memory_efficient_attention(self, *args, **kwargs):
if hasattr(self.prior, 'enable_xformers_memory_efficient_attention'):
self.prior.enable_xformers_memory_efficient_attention(*args, **kwargs)
self.main.enable_xformers_memory_efficient_attention(*args, **kwargs)
def __call__(self, *args, **kwargs):
unclip_outputs = self.prior(prompt=kwargs.get("prompt"), negative_prompt=kwargs.get("negative_prompt"))
if self.prior.device.type == "cuda" or self.prior.device.type == "xpu" or self.prior.device.type == "mps":
prior_device = self.prior.device
self.prior.to("cpu")
self.main.to(prior_device)
kwargs = {**kwargs, **unclip_outputs}
result = self.main(*args, **kwargs)
if self.main.device.type == "cuda" or self.main.device.type == "xpu" or self.prior.device.type == "mps":
main_device = self.main.device
self.main.to("cpu")
self.prior.to(main_device)
return result
def change_backend():
shared.log.info(f'Pipeline changed: {shared.backend}')
unload_model_weights()
@@ -591,6 +544,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
"safety_checker": None,
"requires_safety_checker": False,
"load_safety_checker": False,
"load_connected_pipeline": True # always load end-to-end / connected pipelines
# "use_safetensors": True, # TODO(PVP) - we can't enable this for all checkpoints just yet
}
if devices.dtype == torch.float16:
@@ -699,13 +653,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
if hasattr(sd_model, "watermark"):
sd_model.watermark = NoWatermark()
# Prior pipelines
if hasattr(checkpoint_info, 'model_info') and checkpoint_info.model_info is not None and "prior" in checkpoint_info.model_info:
prior_id = checkpoint_info.model_info["prior"]
shared.log.info(f"Loading diffuser prior: {checkpoint_info.filename} {prior_id}")
prior = diffusers.DiffusionPipeline.from_pretrained(prior_id, **diffusers_load_config)
sd_model = PriorPipeline(prior=prior, main=sd_model) # wrap sd_model
if hasattr(sd_model, "enable_model_cpu_offload"):
if shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload:
shared.log.debug(f'Diffusers {op}: enable model CPU offload')
@@ -828,71 +775,48 @@ class DiffusersTaskType(Enum):
INPAINTING = 3
def set_diffuser_pipe(pipe, new_pipe_type):
wrapper_pipe = None
sd_checkpoint_info = pipe.sd_checkpoint_info
sd_model_checkpoint = pipe.sd_model_checkpoint
sd_model_hash = pipe.sd_model_hash
if pipe.__class__ == PriorPipeline:
wrapper_pipe = pipe
pipe = pipe.main
pipe_name = pipe.__class__.__name__
pipe_name = pipe_name.replace("Img2Img", "").replace("Inpaint", "")
new_pipe_cls_str = None
if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE:
new_pipe_cls_str = pipe_name
new_pipe = diffusers.AutoPipelineForText2Image.from_pipe(pipe)
elif new_pipe_type == DiffusersTaskType.IMAGE_2_IMAGE:
tmp_pipe_name = pipe_name.replace("Pipeline", "Img2ImgPipeline")
if hasattr(diffusers, tmp_pipe_name):
new_pipe_cls_str = pipe_name.replace("Pipeline", "Img2ImgPipeline")
new_pipe = diffusers.AutoPipelineForImage2Image.from_pipe(pipe)
elif new_pipe_type == DiffusersTaskType.INPAINTING:
tmp_pipe_name = pipe_name.replace("Pipeline", "InpaintPipeline")
if hasattr(diffusers, tmp_pipe_name):
new_pipe_cls_str = pipe_name.replace("Pipeline", "InpaintPipeline")
new_pipe = diffusers.AutoPipelineForInpainting.from_pipe(pipe)
if new_pipe_cls_str is None:
shared.log.warning(f'Diffusers unknown pipeline: {tmp_pipe_name}')
new_pipe_cls_str = pipe_name
new_pipe_cls = getattr(diffusers, new_pipe_cls_str)
if pipe.__class__ == new_pipe_cls:
if pipe.__class__ == new_pipe.__class__:
return
new_pipe = new_pipe_cls(**pipe.components)
if wrapper_pipe is not None:
wrapper_pipe.main = new_pipe
new_pipe = wrapper_pipe
new_pipe.sd_checkpoint_info = sd_checkpoint_info
new_pipe.sd_model_checkpoint = sd_model_checkpoint
new_pipe.sd_model_hash = sd_model_hash
model_data.sd_model = new_pipe
shared.log.info(f"Pipeline class changed from {pipe.__class__.__name__} to {new_pipe_cls.__name__}")
shared.log.info(f"Pipeline class changed from {pipe.__class__.__name__} to {new_pipe.__class__.__name__}")
def get_native(pipe: diffusers.DiffusionPipeline):
if pipe.__class__ == PriorPipeline:
pipe = pipe.main
try:
if hasattr(pipe, "vae") and hasattr(pipe.vae.config, "sample_size"):
# Stable Diffusion
size = pipe.vae.config.sample_size
except Exception:
elif hasattr(pipe, "movq") and hasattr(pipe.movq.config, "sample_size"):
# Kandinsky
size = pipe.movq.config.sample_size
elif hasattr(pipe, "unet") and hasattr(pipe.unet.config, "sample_size"):
size = pipe.unet.config.sample_size
else:
size = 0
return size
def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
if pipe.__class__ == PriorPipeline:
pipe = pipe.main
if "Img2Img" in pipe.__class__.__name__:
if pipe.__class__ in diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.values():
return DiffusersTaskType.IMAGE_2_IMAGE
elif "Inpaint" in pipe.__class__.__name__:
elif pipe.__class__ in diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING.values():
return DiffusersTaskType.INPAINTING
return DiffusersTaskType.TEXT_2_IMAGE
@@ -1072,10 +996,6 @@ def apply_token_merging(sd_model, token_merging_ratio):
if current_token_merging_ratio > 0:
tomesd.remove_patch(sd_model)
if sd_model.__class__ == PriorPipeline:
# token merging is not supported for PriorPipelines currently
return
if token_merging_ratio > 0:
tomesd.apply_patch(
sd_model,
+1 -1
View File
@@ -49,7 +49,7 @@ requests==2.31.0
tqdm==4.65.0
accelerate==0.20.3
opencv-python==4.7.0.72
diffusers==0.18.2
diffusers==0.19.0
einops==0.4.1
gradio==3.32.0
numexpr==2.8.4