From ff28ac35e044052d1d90d7b05616bc33d23c9401 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 25 Sep 2023 11:23:04 -0400 Subject: [PATCH] add actual latent upscalers --- CHANGELOG.md | 11 + cli/hf-search.py | 2 +- javascript/style.css | 1 + modules/dml/hijack/realesrgan_model.py | 2 +- modules/freeu/freeu_diffusers.py | 306 +++++++++++++++++++++++ modules/modelloader.py | 2 +- modules/paths.py | 35 ++- modules/postprocess/sdupscaler4_model.py | 64 +++++ modules/processing_diffusers.py | 9 +- modules/sd_models.py | 151 ++++++----- modules/ui_models.py | 2 +- 11 files changed, 494 insertions(+), 91 deletions(-) create mode 100644 modules/freeu/freeu_diffusers.py create mode 100644 modules/postprocess/sdupscaler4_model.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 885e6b471..bdacba68d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,8 +21,19 @@ Upgrades are still possible and supported, but above is recommended for best exp - faster search, ability to show/hide/sort networks - refactored subfolder handling *note*: this will trigger model hash recaclulation on first model use +- **Refiner**: + - You can now use *SD Latent Upscale* models as refiner + this is a bit experimental, but it works quite well! + Simply go to *Models -> Huggingface* and download: + - `stabilityai/sd-x2-latent-upscaler` + - `stabilityai/stable-diffusion-x4-upscaler` - **Upscalers**: - more high quality upscalers available by default + *SwinIR:2, ESRGAN:12, RealESRGAN:6, SCUNet:2* + - two additional latent upscalers based on SD upscale models when using Diffusers backend + *SD Upscale 2x, SD Upscale 4x* + Note: Recommended usage for *SD Upscale* is by using second pass instead of upscaler + as it allows for tuning of prompt, seed, sampler settings which are used to guide upscaler - unified init/download/execute/progress code - easier installation - available in **xyz grid** diff --git a/cli/hf-search.py b/cli/hf-search.py index ac97b6c26..53c254d44 100755 --- a/cli/hf-search.py +++ b/cli/hf-search.py @@ -10,7 +10,7 @@ if __name__ == "__main__": hf_api = hf.HfApi() model_filter = hf.ModelFilter( model_name=keyword, - task='text-to-image', + # task='text-to-image', library=['diffusers'], ) res = hf_api.list_models(filter=model_filter, full=True, limit=50, sort="downloads", direction=-1) diff --git a/javascript/style.css b/javascript/style.css index 0e3cd4f62..62214f207 100644 --- a/javascript/style.css +++ b/javascript/style.css @@ -31,6 +31,7 @@ div.gradio-html.min{ min-height: 0; } .settings-accordion .gap { padding-right: 1000px; } .small-accordion { width: fit-content !important; padding-left: 0 !important; } .small-accordion .form { min-width: var(--left-column) !important; } +.small-accordion .label-wrap { padding: 16px 0px 8px 0px; margin: 0; border-top: 2px solid var(--button-secondary-border-color); } .small-accordion .label-wrap .icon { margin-right: 1.6em; margin-left: 0.6em; color: var(--button-primary-border-color); } .hidden { display: none; } footer { display: none; } diff --git a/modules/dml/hijack/realesrgan_model.py b/modules/dml/hijack/realesrgan_model.py index b55e14647..bee137a45 100644 --- a/modules/dml/hijack/realesrgan_model.py +++ b/modules/dml/hijack/realesrgan_model.py @@ -1,6 +1,6 @@ import math import torch -from modules.realesrgan_model_arch import RealESRGANer +from modules.postprocess.realesrgan_model_arch import RealESRGANer # DML Solution: Some of contents of output tensor turn to 0 after Extended Slices. Move it to cpu. diff --git a/modules/freeu/freeu_diffusers.py b/modules/freeu/freeu_diffusers.py new file mode 100644 index 000000000..a9a82f1e6 --- /dev/null +++ b/modules/freeu/freeu_diffusers.py @@ -0,0 +1,306 @@ +# https://github.com/lyn-rgb/FreeU_Diffusers/blob/diffusers-v0.21.2/free_lunch_utils.py + +""" +register_free_upblock2d(pipe) +register_free_crossattn_upblock2d(pipe) +""" + + +from typing import Any, Dict, Optional, Tuple +import torch +import torch.fft as fft +from diffusers.utils import is_torch_version + + +def isinstance_str(x: object, cls_name: str): + """ + Checks whether x has any class *named* cls_name in its ancestry. + Doesn't require access to the class's implementation. + + Useful for patching! + """ + + for _cls in x.__class__.__mro__: + if _cls.__name__ == cls_name: + return True + + return False + + +def Fourier_filter(x, threshold, scale): + dtype = x.dtype + B, C, H, W = x.shape + # Non-power of 2 images must be float32 + if (W & (W - 1)) != 0 or (H & (H - 1)) != 0: + x = x.type(torch.float32) + # FFT + x_freq = fft.fftn(x, dim=(-2, -1)) + x_freq = fft.fftshift(x_freq, dim=(-2, -1)) + + B, C, H, W = x_freq.shape + mask = torch.ones((B, C, H, W)).to(x.device) + + crow, ccol = H // 2, W //2 + mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale + x_freq = x_freq * mask + + # IFFT + x_freq = fft.ifftshift(x_freq, dim=(-2, -1)) + x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real + + x_filtered = x_filtered.type(dtype) + return x_filtered + + +def register_upblock2d(model): + def up_forward(self): + def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states + + return forward + + for _i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "UpBlock2D"): + upsample_block.forward = up_forward(upsample_block) + + +def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + def up_forward(self): + def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + #print(f"in free upblock2d, hidden states shape: {hidden_states.shape}") + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # --------------- FreeU code ----------------------- + # Only operate on the first two stages + if hidden_states.shape[1] == 1280: + hidden_states[:,:640] = hidden_states[:,:640] * self.b1 + res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1) + if hidden_states.shape[1] == 640: + hidden_states[:,:320] = hidden_states[:,:320] * self.b2 + res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2) + # --------------------------------------------------------- + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states + + return forward + + for _i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "UpBlock2D"): + upsample_block.forward = up_forward(upsample_block) + upsample_block.b1 = b1 + upsample_block.b2 = b2 + upsample_block.s1 = s1 + upsample_block.s2 = s2 + + +def register_crossattn_upblock2d(model): + def up_forward(self): + def forward( + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + return hidden_states + + return forward + + for _i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "CrossAttnUpBlock2D"): + upsample_block.forward = up_forward(upsample_block) + + +def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + def up_forward(self): + def forward( + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + #print(f"in free crossatten upblock2d, hidden states shape: {hidden_states.shape}") + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # --------------- FreeU code ----------------------- + # Only operate on the first two stages + if hidden_states.shape[1] == 1280: + hidden_states[:,:640] = hidden_states[:,:640] * self.b1 + res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1) + if hidden_states.shape[1] == 640: + hidden_states[:,:320] = hidden_states[:,:320] * self.b2 + res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2) + # --------------------------------------------------------- + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + return hidden_states + + return forward + + for _i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "CrossAttnUpBlock2D"): + upsample_block.forward = up_forward(upsample_block) + upsample_block.b1 = b1 + upsample_block.b2 = b2 + upsample_block.s1 = s1 + upsample_block.s2 = s2 diff --git a/modules/modelloader.py b/modules/modelloader.py index 3c4b93934..d96392d56 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -249,7 +249,7 @@ def find_diffuser(name: str): hf_api = hf.HfApi() hf_filter = hf.ModelFilter( model_name=name, - task='text-to-image', + # task='text-to-image', library=['diffusers'], ) models = list(hf_api.list_models(filter=hf_filter, full=True, limit=20, sort="downloads", direction=-1)) diff --git a/modules/paths.py b/modules/paths.py index d8791776a..6baed946c 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -50,31 +50,30 @@ def create_paths(opts, log=None): def create_path(folder): if folder is None or folder == '': return - if not os.path.exists(folder): - try: - os.makedirs(folder, exist_ok=True) - if log is not None: - log.debug(f'Create path: {folder}') - except Exception as e: - if log is not None: - log.error(f'Failed to create path: {folder} {e}') + if os.path.exists(folder): + return + try: + os.makedirs(folder, exist_ok=True) + if log is not None: + log.debug(f'Create folder={folder}') + except Exception as e: + if log is not None: + log.error(f'Create Failed folder={folder} {e}') def fix_path(folder): tgt = opts.data.get(folder, None) or opts.data_labels[folder].default if tgt is None or tgt == '': return tgt + if os.path.isabs(tgt): + return tgt if len(data_path) > 0 and tgt.startswith(data_path): # path is already relative to data_path return tgt - fullpath = os.path.join(data_path, tgt) - if len(data_path) > 0 and os.path.isabs(data_path): - return fullpath - if os.path.isabs(fullpath) and os.path.exists(fullpath): - return fullpath - try: - relpath = os.path.relpath(fullpath, script_path) - opts.data[folder] = relpath - except Exception: - opts.data[folder] = fullpath + else: + tgt = os.path.join(data_path, tgt) + if os.path.isabs(tgt): + return tgt + tgt = os.path.relpath(tgt, script_path) + opts.data[folder] = tgt return opts.data[folder] create_path(data_path) diff --git a/modules/postprocess/sdupscaler4_model.py b/modules/postprocess/sdupscaler4_model.py new file mode 100644 index 000000000..574f02cf1 --- /dev/null +++ b/modules/postprocess/sdupscaler4_model.py @@ -0,0 +1,64 @@ +import torch +import diffusers +from PIL import Image +from modules import shared, devices +from modules.upscaler import Upscaler, UpscalerData + +class UpscalerSD(Upscaler): + def __init__(self, dirname): # pylint: disable=super-init-not-called + self.name = "StableDiffusion" + self.user_path = dirname + if shared.backend != shared.Backend.DIFFUSERS: + super().__init__() + return + self.scalers = [ + UpscalerData(name="SD Latent 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), + UpscalerData(name="SD Latent 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), + ] + self.pipelines = [ + None, + None, + ] + + def load_model(self, path: str): + from modules.sd_models import set_diffuser_options + scaler = [x for x in self.scalers if x.data_path == path][0] + if scaler.model is None: + devices.set_cuda_params() + scaler.model = diffusers.DiffusionPipeline.from_pretrained(path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype) + if hasattr(scaler.model, "set_progress_bar_config"): + scaler.model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + 'Upscale', ncols=80, colour='#327fba') + set_diffuser_options(scaler.model, vae=None, op='upscaler') + return scaler.model + + def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor): + pass + + def do_upscale(self, img: Image.Image, selected_model): + devices.torch_gc() + model = self.load_model(selected_model) + if model is None: + return img + seeds = [torch.randint(0, 2 ** 32, (1,)).item() for _ in range(1)] + generator_device = devices.cpu if shared.opts.diffusers_generator_device == "cpu" else devices.device + generator = [torch.Generator(generator_device).manual_seed(s) for s in seeds] + args = { + 'prompt': '', + 'negative_prompt': '', + 'image': img, + 'num_inference_steps': 20, + 'guidance_scale': 7.5, + 'generator': generator, + 'latents': None, + 'return_dict': True, + 'callback': self.callback, + 'callback_steps': 1, + # 'noise_level': 100, + # 'num_images_per_prompt': 1, + # 'eta': 0.0, + # 'cross_attention_kwargs': None, + } + model = model.to(devices.device) + output = model(**args) + image = output.images[0] + return image diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index c382697ef..b099fe195 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -443,6 +443,12 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro # shared.log.warning(f'Refiner requires image size to be divisible by 8: {image.shape}') # results.append(image) # return results + noise_level = round(350 * p.denoising_strength) + output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np', + if shared.sd_refiner.__class__.__name__ == 'StableDiffusionUpscalePipeline': + image = vae_decode(latents=image, model=shared.sd_model, full_quality=p.full_quality, output_type='pil') + p.extra_generation_params['Noise level'] = noise_level + output_type = 'np' refiner_args = set_pipeline_args( model=shared.sd_refiner, prompts=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts[i], @@ -450,12 +456,13 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro num_inference_steps=int(p.refiner_steps // (1 - p.refiner_start)) if p.refiner_start > 0 and p.refiner_start < 1 and refiner_is_sdxl else int(p.refiner_steps // p.denoising_strength + 1) if refiner_is_sdxl else p.refiner_steps, eta=shared.opts.scheduler_eta, strength=p.denoising_strength, + noise_level=noise_level, # StableDiffusionUpscalePipeline only guidance_scale=p.image_cfg_scale if p.image_cfg_scale is not None else p.cfg_scale, guidance_rescale=p.diffusers_guidance_rescale, denoising_start=p.refiner_start if p.refiner_start > 0 and p.refiner_start < 1 else None, denoising_end=1 if p.refiner_start > 0 and p.refiner_start < 1 else None, image=image, - output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np', + output_type=output_type, clip_skip=p.clip_skip, desc='Refiner', ) diff --git a/modules/sd_models.py b/modules/sd_models.py index 1c1e51e48..ed9934d20 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -689,6 +689,76 @@ def compile_diffusers(sd_model): except Exception as err: shared.log.warning(f"Model compile not supported: {err}") + +def set_diffuser_options(sd_model, vae, op: str): + if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram): + shared.log.warning(f'Setting {op}: Model CPU offload and Sequential CPU offload are not compatible') + shared.log.debug(f'Setting {op}: disabling model CPU offload') + shared.opts.diffusers_model_cpu_offload=False + shared.cmd_opts.medvram=False + + if hasattr(sd_model, "watermark"): + sd_model.watermark = NoWatermark() + sd_model.has_accelerate = False + if hasattr(sd_model, "enable_model_cpu_offload"): + if (shared.cmd_opts.medvram and devices.backend != "directml") or shared.opts.diffusers_model_cpu_offload: + shared.log.debug(f'Setting {op}: enable model CPU offload') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') + sd_model.enable_model_cpu_offload() + sd_model.has_accelerate = True + if hasattr(sd_model, "enable_sequential_cpu_offload"): + if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload: + shared.log.debug(f'Setting {op}: enable sequential CPU offload') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') + sd_model.enable_sequential_cpu_offload(device=devices.device) + sd_model.has_accelerate = True + if hasattr(sd_model, "enable_vae_slicing"): + if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_slicing: + shared.log.debug(f'Setting {op}: enable VAE slicing') + sd_model.enable_vae_slicing() + else: + sd_model.disable_vae_slicing() + if hasattr(sd_model, "enable_vae_tiling"): + if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_tiling: + shared.log.debug(f'Setting {op}: enable VAE tiling') + sd_model.enable_vae_tiling() + else: + sd_model.disable_vae_tiling() + if hasattr(sd_model, "enable_attention_slicing"): + if shared.cmd_opts.lowvram or shared.opts.diffusers_attention_slicing: + shared.log.debug(f'Setting {op}: enable attention slicing') + sd_model.enable_attention_slicing() + else: + sd_model.disable_attention_slicing() + if hasattr(sd_model, "vae"): + if vae is not None: + sd_model.vae = vae + if shared.opts.diffusers_vae_upcast != 'default': + if shared.opts.diffusers_vae_upcast == 'true': + # sd_model.vae.config["force_upcast"] = True + sd_model.vae.config.force_upcast = True + else: + # sd_model.vae.config["force_upcast"] = False + sd_model.vae.config.force_upcast = False + if shared.opts.no_half_vae: + devices.dtype_vae = torch.float32 + sd_model.vae.to(devices.dtype_vae) + shared.log.debug(f'Setting {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}') + if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'): + sd_model.enable_xformers_memory_efficient_attention() + if shared.opts.opt_channelslast: + shared.log.debug(f'Setting {op}: enable channels last') + sd_model.unet.to(memory_format=torch.channels_last) + + def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model'): # pylint: disable=unused-argument import torch # pylint: disable=reimported,redefined-outer-name if timer is None: @@ -754,11 +824,21 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No diffusers_load_config["vae"] = vae if os.path.isdir(checkpoint_info.path): - try: + err1 = None + err2 = None + try: # try autopipeline first sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) sd_model.model_type = sd_model.__class__.__name__ except Exception as e: - shared.log.error(f'Failed loading {op}: {checkpoint_info.path} {e}') + err1 = e + try: # try diffusion pipeline next + if err1 is not None: + sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model.model_type = sd_model.__class__.__name__ + except Exception as e: + err2 = e + if err2 is not None: + shared.log.error(f'Failed loading {op}: {checkpoint_info.path} autopipeline={err1} diffusionpipeline={err2}') return elif os.path.isfile(checkpoint_info.path) and checkpoint_info.path.lower().endswith('.safetensors'): diffusers_load_config["local_files_only"] = True @@ -805,72 +885,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No elif "Kandinsky" in sd_model.__class__.__name__: sd_model.scheduler.name = 'DDIM' - if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram): - shared.log.warning(f'Setting {op}: Model CPU offload and Sequential CPU offload are not compatible') - shared.log.debug(f'Setting {op}: disabling model CPU offload') - shared.opts.diffusers_model_cpu_offload=False - shared.cmd_opts.medvram=False - - if hasattr(sd_model, "watermark"): - sd_model.watermark = NoWatermark() - sd_model.has_accelerate = False - if hasattr(sd_model, "enable_model_cpu_offload"): - if (shared.cmd_opts.medvram and devices.backend != "directml") or shared.opts.diffusers_model_cpu_offload: - shared.log.debug(f'Setting {op}: enable model CPU offload') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') - sd_model.enable_model_cpu_offload() - sd_model.has_accelerate = True - if hasattr(sd_model, "enable_sequential_cpu_offload"): - if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload: - shared.log.debug(f'Setting {op}: enable sequential CPU offload') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') - sd_model.enable_sequential_cpu_offload(device=devices.device) - sd_model.has_accelerate = True - if hasattr(sd_model, "enable_vae_slicing"): - if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_slicing: - shared.log.debug(f'Setting {op}: enable VAE slicing') - sd_model.enable_vae_slicing() - else: - sd_model.disable_vae_slicing() - if hasattr(sd_model, "enable_vae_tiling"): - if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_tiling: - shared.log.debug(f'Setting {op}: enable VAE tiling') - sd_model.enable_vae_tiling() - else: - sd_model.disable_vae_tiling() - if hasattr(sd_model, "enable_attention_slicing"): - if shared.cmd_opts.lowvram or shared.opts.diffusers_attention_slicing: - shared.log.debug(f'Setting {op}: enable attention slicing') - sd_model.enable_attention_slicing() - else: - sd_model.disable_attention_slicing() - if hasattr(sd_model, "vae"): - if vae is not None: - sd_model.vae = vae - if shared.opts.diffusers_vae_upcast != 'default': - if shared.opts.diffusers_vae_upcast == 'true': - # sd_model.vae.config["force_upcast"] = True - sd_model.vae.config.force_upcast = True - else: - # sd_model.vae.config["force_upcast"] = False - sd_model.vae.config.force_upcast = False - if shared.opts.no_half_vae: - devices.dtype_vae = torch.float32 - sd_model.vae.to(devices.dtype_vae) - shared.log.debug(f'Setting {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}') - if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'): - sd_model.enable_xformers_memory_efficient_attention() - if shared.opts.opt_channelslast: - shared.log.debug(f'Setting {op}: enable channels last') - sd_model.unet.to(memory_format=torch.channels_last) + set_diffuser_options(sd_model, vae, op) base_sent_to_cpu=False if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none') or shared.opts.ipex_optimize: diff --git a/modules/ui_models.py b/modules/ui_models.py index 182f45536..e2bb93d74 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -187,7 +187,7 @@ def create_ui(): hf_api = hf.HfApi() model_filter = hf.ModelFilter( model_name=keyword, - task='text-to-image', + # task='text-to-image', library=['diffusers'], ) models = hf_api.list_models(filter=model_filter, full=True, limit=50, sort="downloads", direction=-1)