Add SUPIR Upscaler - Githubissues

DN6 commented 7 months ago

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

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

Provide useful links for the implementation

No response

nxbringr commented 7 months ago

Hey @DN6, can I please work on this?

yiyixuxu commented 7 months ago

@ihkap11 hey! sure!

Bhavay-2001 commented 6 months ago

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

nxbringr commented 6 months ago

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 commented 6 months ago

ok great. Pls let me know. Thanks

landmann commented 6 months ago

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

nxbringr commented 6 months ago

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 commented 6 months ago

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

nxbringr commented 6 months ago

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 commented 6 months ago

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 commented 6 months ago

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

nxbringr commented 6 months ago

Hey @landmann, I'm finding it hard to map a few components from the paper's network architecture details to the codebase they provided.

Currently, I'm stuck on understanding how they are trimming the VIT blocks when using the modified ControlNet adapter with ZeroFST connector at the code level. They seem to use GLVControl but no VIT component and network trimming that I can spot in the codebase.

I sent an email to one of the authors last week. If I don't hear back, I plan to follow up with more specific questions this week. (Also check this issue here) I'm playing with the code in my repo atm here

If interested would you want to take a second look at their code and share your thoughts?

landmann commented 6 months ago

@ihkap11 are you following the paper and trying to code it? Why not just make a wrapper around what they have? It's pytorch after all, no? I haven't read the paper much in depth, but was able to run SUPIR locally!

@austinfujimori you should take a look if you're free 🙂