add actual latent upscalers

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**  

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)

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; }

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.

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

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))

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)

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

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',
             )

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:

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)