Kohya Hires fix

[crapthings]

is diffusers possible to support this hires fix?

it looks 1.5 work too

AUTOMATIC1111/stable-diffusion-webui#13974

https://www.youtube.com/watch?v=SbgMwHDXthU

same seed at 1024x1024

with

without

This method will also improve the quality of inpainting without two pass

[crapthings]

without

with

[sayakpaul]

I think the improvements are quite amazing TBH. Ccing @kohya-ss in case you would like to elaborate the motivation behind the fix a bit.

@yiyixuxu WDYT about shipping this?

[crapthings]

this one is 1280 x 1280 with Kohya Hires.fix

[sayakpaul]

Gently pinging @kohya-ss since you reacted to my comment. If you find a moment, would love some sort of a brief explanation about the feature. Would be really helpful!

[kohya-ss]

Hi!

This method is extremely simple.

The distortion of composition at large resolutions can be attributed to the U-Net's limited capacity to handle feature maps of such sizes effectively.

In Stable Diffusion, it has been observed that the composition is primarily determined in the deeper parts of the U-Net, as evident from the effects of block merging explored by the community.

Furthermore, the previews of the denoising steps reveal that the composition is also heavily influenced by the timesteps closer to the noise (i.e., the larger timesteps) during the denoising process.

Based on these observations, this method proposes resizing the feature maps to a smaller size specifically at the timesteps close to the noise in the deeper parts of the U-Net. This allows the U-Net to process the feature maps more effectively, thereby mitigating the distortion of the composition.

Please refer to the following link for the specific implementation code. In the input blocks of the U-Net at certain depths, the hidden states are downsampled, and in the corresponding depths of the output blocks, they are upsampled.

https://github.com/kohya-ss/sd-scripts/blob/2d7389185c021bc527b414563c245c5489d6328a/library/sdxl_original_unet.py#L1192

In this implementation, two depths (ds_depth1 and ds_depth2) can be specified to apply the method, which remain effective up to the timesteps ds_timesteps_1 and ds_timesteps_2, respectively. ds_ratio represents the shrinking ratio.

Additionally, laksjdjf has created a ComfyUI node for this method, which can be found at the following link:
https://gist.github.com/laksjdjf/487a28ceda7f0853094933d2e138e3c6

Furthermore, although I have not personally used it, ComfyUI seems to officially support this method. Their node allows users to specify arbitrary depths, timesteps, and shrinking ratios, and it also appears to support cascading multiple instances of the node.

[sayakpaul]

Thanks much for your detailed explanation. Really appreciate it.

[jorgemcgomes]

Looks very interesting.
Could this theoretically be used in the opposite way, to generate smaller images?
When trying to generate small images with models that have been trained with high-res images, we often see the same kind of problems -- weird compositions, low quality.
So, could the feature maps be resized to a larger size? So that it becomes possible to generate low-res images that have the same composition as high-res ones?

[blx0102]

looking forward to this feature!

[yiyixuxu]

can someone make a community pipeline first for this? basically, you need to:

1. copy over the UNet2DConditionModel, rename it, add slightly modify the forward method to implement kohya Hires fix
2. add a from_unet method which is basically

@classmethod
    def from_unet(
        cls,
        unet: UNet2DConditionModel,
        ds_depth_1 = ds_depth_1,
        ds_depth_2 = ds_depth_2,
        ...
    ):  
        config = dict((unet.config)
        config[" ds_depth_1"] = ds_depth_1
        config["ds_depth_2"] = ds_depth_2,

        koyha_unet = cls.from_config(config)
        koyha_unet.load_state_dict(unet.state_dict())

3. copy over the SD or SDXL pipeline, rename it, and slightly update the __init__ method

def __init__():
      ...
       unet = kohyaHiresFixUnet2DConditionalModel.from_unet(unet, **kwargs)
        self.register_modules(
           ...
            unet=unet,
        )

and then please play around with it, figure out how to configure it for SD and SDXL, and show us some results! Once we have the community pipeline that we can play with, depending on how good the results are and how users react to it, we can then decide if we want to integrate this fix into our Unet or not (should be pretty straightforward)

[sajadn]

Hey, I'll work on this.