huggingface
Add SUPIR Upscaler
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
Hi @yiyixuxu, anyone working on this? Can I also contribute? Please let me know how may I proceed?
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 :)
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
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:
-
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.
-
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.
-
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.
- 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.
- 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.
-
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.
-
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.safetensorsas 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 :).
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
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?
@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 🙂
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?
Hi, I tried don't load the VIT ckpt, and it has no influence !!!
@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 🙂
the "ZeroSFT" is to replace the concat[hidden_states, res_sample] in AttnUpBlock and CrossAttnUpBlock, so we can't use diffusers's sdxl pipeline to implement it.
@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 🙂
the "ZeroSFT" is to replace the concat[hidden_states, res_sample] in AttnUpBlock and CrossAttnUpBlock, so we can't use diffusers's sdxl pipeline to implement it.
because the different between the sgm and diffusers arch, its difficult
@CuddleSabe would you like to connect on discord to discuss SUPIR and possibly collaborate in figuring out how to support it in diffusers? You can find me by the username tortillachips11.
@CuddleSabe would you like to connect on discord to discuss SUPIR and possibly collaborate in figuring out how to support it in diffusers? You can find me by the username tortillachips11.
well, cant use discord. I write a train script and the model file to train sd1.5 supir. however, I cant publish it because the Company Confidentiality Law
@CuddleSabe would you like to connect on discord to discuss SUPIR and possibly collaborate in figuring out how to support it in diffusers? You can find me by the username tortillachips11.
in sgm, it like this
the "hs" is the res_sample from the down_blocks. https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/openaimodel.py#L849
but in diffusers, it like this
the "res_sample" in diffusers equal the "hs" in sgm
https://github.com/huggingface/diffusers/blob/66f94eaa0c68a893b2aba1ec9f79ee7890786fba/src/diffusers/models/unets/unet_2d_condition.py#L1189
so you need to rewrite the AttnUpBlock2D 、CrossAttnUpBlock2D and UNetMidBlock2DCrossAttn in diffusers
Hey! I tried but couldn't get this working. Feel free to take over the implementation for this Issue.
Hey! I tried but couldn't get this working. Feel free to take over the implementation for this Issue.
But do you have a branch where we can continue where you left off? I might try this after I finish a project I'm involved with.
Just in case, this is not an easy task, everything is in the sgm format so there's a lot of conversion involved. It requires a deep understanding of the original code and diffusers.
Probably the best choice here is to start as a research project and convert all the sgm code to diffusers, and then when stuck, get help from the maintainers and the community.
accoding to the paper and the Comfyui will impliment below:
- SUPIR MODEL_LOADER -> SUPIR_MODEL, SUPIR_vae
- SUPIR_VAE = vae.from_config and load conveted state_dict
--
- SUPIR_FIRSTSTAGE Denoiser : take Low quality image in and blur or smooth image out and it`s latent
- this stage include an SUPIR VAE include vae-encoder vae-decoder
- SUPIR_VAE.encoder(LQ_image) ->supir_latent -> SUPIR_VAE.decoder (supir_latent).
3 SUPIR-controlnet : which take latents in and time stpes in , generate controlnet residuals dowsamples and midsample out
- Class trim_controlnet(Controlnet_normal):
not impliment
4 An hacked Unet which modify the connector of each dow and up blocks use ZeroSFT
- replace Unet-zero_conv-connector
hi @DN6 i already did it and it is working fine see results: https://imgsli.com/Mzg2NTUx/0/1. But then I saw that they changed the license and it seems that it is very restrictive. I don't know how to deal with this case. I even sent an email to the owner but it's been a month since I got a response.
They changed the license from MIT to "something unnamed" literally two days after the original post here, didn't see that.
Don't know why they keep changing it to a custom one, probably because there's some other party involved, but basically, this license is the same as Flux, you can use it and change it freely if it's non-commercial which is fine, we have other models with the same type of custom license.
AFAIK this is still the SOTA upscaler model so it would be nice to have it.
They changed the license from MIT to "something unnamed" literally two days after the original post here, didn't see that.
Don't know why they keep changing it to a custom one, probably because there's some other party involved, but basically, this license is the same as Flux, you can use it and change it freely if it's non-commercial which is fine, we have other models with the same type of custom license.
AFAIK this is still the SOTA upscaler model so it would be nice to have it.
they have this on their new license: "b. Proprietary Neural Network Components: Notwithstanding the open-source license granted in Section 2(a) for open-source components, the Licensor retains all rights, title, and interest in and to the proprietary Neural Network Components developed and trained by the Licensor. The license granted herein does not include rights to use, copy, modify, distribute, or create derivative works of these proprietary Neural Network Components without obtaining express written permission from the Licensor."
The license is a bit contradictory because at the same time it says that parts of the code are open source, another part cannot be changed. I wonder if in the case where I converted the model to .safetensors I would be violating this clause.
I understand that part (I'm not expert) as to referring to derivative works only, that's a clause for limiting the use of the model to do something similar to circumvent the license and then use it commercially, like using its outputs to train a similar model.
When integrating it with diffusers and changing the format for diffusers, and of course copying the license and referencing them as copyright owners, you're not using it as a derivative work but as the original model.
On the other part, if we interpret it as for everything and not just derivative works, it invalidates all from before and at that point they should just remove "open source" from it and just say that any use of the model will need their permission.
But still it would be a lot better to get the permission directly from the authors, since it's a not easy to understand license.