Add SUPIR Upscaler

[DN6]

Model/Pipeline/Scheduler description

SUPIR is a super-resolution model that looks like it produces excellent results

Github Repo: https://github.com/Fanghua-Yu/SUPIR

The model is quite memory intensive, so the optimisation features available in diffusers might be quite helpful in making this accessible to lower resource GPUs.

Open source status

• The model implementation is available.
• The model weights are available (Only relevant if addition is not a scheduler).

Provide useful links for the implementation

No response

[nxbringr]

Hey @DN6, can I please work on this?

[yiyixuxu]

@ihkap11 hey! sure!

[Bhavay-2001]

Hi @yiyixuxu, anyone working on this? Can I also contribute? Please let me know how may I proceed?

[nxbringr]

Hey @Bhavay-2001 I'm currently working on this. Will post the PR here soon.
I can tag you on the PR if I there is something I need help with :)

[Bhavay-2001]

ok great. Pls let me know.
Thanks

[landmann]

@ihkap11 how's it going 😁 I'd loooooove to have this

[nxbringr]

Hey @landmann I'll post the PR this weekend and tag you if you want to contribute to it :) apologies for the delay, it's my first new model implementation PR

[landmann]

You a real champ 🙌
Happy Friday, my gal/dude!

[nxbringr]

Initial Update:

• Understood the paper (paper highlights below)
• Currently defining paper components that will become diffuser artefacts. (WIP: breaking down SUPIR code)

Paper Insights

Motivation:

• IR methods based on generative priors leverage powerful pre-trained generative models to introduce high-quality generation and prior knowledge into IR, bringing significant progress in
  perceptual effects and intelligence of IR results.
• Continuously enhancing the capabilities of the generative prior is key to achieving more intelligent IR results, with model scaling being a crucial and effective approach.
• The authors propose scaling up generative priors and training data to address these limitations.

Architecture Overview:

1. Generative Prior: The authors choose SDXL (Stable Diffusion XL) as the backbone for their generative prior due to its high-resolution image generation capability without hierarchical design.

2. Degradation-Robust Encoder: They fine-tune the SDXL encoder to make it robust to degradation, enabling effective mapping of low-quality (LQ) images to the latent space.

3. Large-Scale Adaptor: The author designed a new adaptor with network trimming and a ZeroSFT connector to control the generation process at the pixel level.

   Issues with existing adaptors

   
   

   • LoRA limits generation but struggles with LQ image control
   • T2I lacks the capacity for effective LQ image content identification
   • ControlNet’s direct copy is challenging for the SDXL model scale.
   1. Network Trimming: Modify the adaptor architecture by trimming half of the ViT blocks in each encoder block (of SDXL) to achieve a balance between network capacity and computational feasibility.
   2. Redesigning the Connector: The introduced ZeroSFT module is built upon zero convolution and incorporates an additional spatial feature transfer (SFT) operation and group normalization.
   Why do we need this?

   
   

   • The authors note that while SDXL's generative capacity delivers excellent visual effects, it also makes precise pixel-level control challenging.
   • ControlNet uses zero convolution for generation guidance, but relying solely on residuals is insufficient for the level of control required by IR tasks.
     

4. Multi-Modality Language Guidance: They incorporate the LLaVA multi-modal large language model to understand image content and guide the restoration process using textual prompts.

5. Restoration-Guided Sampling: They propose a modified sampling method to selectively guide the prediction results to be close to the LQ image, ensuring fidelity in the restored image.

Thoughts on implementation details:

• Trainable components are degradation robust encoder and trimmed ControlNet.
• Extend the SDXL class from Diffusers and use SDXL checkpoint = sd_xl_base_1.0_0.9vae.safetensors as base pre-trained generative prior.
• The SUPIR model will first load pre-trained weights from the SDXL checkpoint, then it will load SUPIR-specific weights, which include the modifications and additions made to adapt the SDXL model for image restoration tasks.
• Trimmed ControlNet encoder which trims half of the ViT blocks from each encoder block. (Todo: Figure out where to make this change)
• In the SUPIR model, SDXL (Stable Diffusion XL) is used as the backbone for the generative prior, and the GLVControl and LightGLVUNet modules are used as the adaptor to guide the SDXL model for image restoration. Todo: Convert to Diffusers Artifact
• Probably, a dummy code would look like this:

class SUPIRModel(nn.Module):
    def __init__(self, sdxl_model_path):
        super().__init__()
        self.sdxl_pipeline = StableDiffusionXLPipeline.from_pretrained(sdxl_model_path)
        self.glv_control = GLVControl(in_channels=3, out_channels=64, context_dim=128)
        self.light_glv_unet = LightGLVUNet(in_channels=3, out_channels=3)
        
    def forward(self, lq_image, context, num_inference_steps=50):
        # Generate control signal using GLVControl
        control_signal = self.glv_control(lq_image, context)
        
        # Use SDXL pipeline for guided diffusion
        restored_image = self.sdxl_pipeline(
            prompt="",
            image=lq_image,
            control_image=control_signal,
            num_inference_steps=num_inference_steps,
            generator=None,
        ).images[0]
        
        # Refine the restored image using LightGLVUNet
        refined_image = self.light_glv_unet(restored_image, control_signal)
        
        return refined_image

• ZeroFST acts as a connector. Todo: Convert to Diffusers Artifact

To cover later:

• LLaVA for multi-modality language guidance.

I'm currently in the process of breaking down SUPIR code into diffusers artefacts and figuring out optimization techniques to make it compatible with low-resource GPUs.

Feel free to correct me or start a discussion on this thread. Let me know if you wish to collaborate, I'm happy to set up discussions and work on it together :).

[landmann]

Looks fantastic! How far along did you get, @ihkap11 ?

Btw, a good reference for the input parameters are here https://replicate.com/cjwbw/supir?prediction=32glqstbvpjjppxmvcge5gsncu

[landmann]

@ihkap11 how you doing? Which part are you stuck?