CCSR-v2 Model Weights

This repository hosts the pre-trained weights for CCSR-v2 (Continuous Latent Diffusion for Image Super-Resolution), optimized for stable and step-flexible image upscaling.

• Official Source Code: AI-Wrappers/ccsr-v2-pruned (GitHub)
• Python Library: ccsr-pruned (PyPI)
• Original Paper: arXiv:2401.00877

📁 Repository Structure

This model repost contains the following folders:

1. stable-diffusion-2-1-base/ - Standard components of the Stable Diffusion 2.1 base model.
2. controlnet/ - Custom ControlNet weights trained for Stage 1 of CCSR.
3. vae/ - Fine-tuned VAE weights trained for Stage 2 of CCSR.

🚀 Quick Start / Inference

To run super-resolution with these weights, you can use the official ccsr-pruned Python library.

1. Install the Library

pip install ccsr-pruned
# or using uv
uv add ccsr-pruned

2. Python Code Example

import torch
from PIL import Image
from diffusers import AutoencoderKL
from ccsr import StableDiffusionControlNetCCSRPipeline, ControlNetCCSRModel

# Load the models using this Hugging Face repository as the source
repo_id = "kharma1/ccsr_v2_repost"

# 1. Load ControlNet, VAE, and Pipeline
controlnet = ControlNetCCSRModel.from_pretrained(repo_id, subfolder="controlnet", torch_dtype=torch.float16)
vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=torch.float16)

pipeline = StableDiffusionControlNetCCSRPipeline.from_pretrained(
    f"{repo_id}/stable-diffusion-2-1-base",
    controlnet=controlnet,
    vae=vae,
    torch_dtype=torch.float16,
).to("cuda")

# 2. Prepare low-quality input image (upscaled and divisible by 8)
lq_image = Image.open("input_lq.png").convert("RGB")
upscale_factor = 4
target_width = lq_image.size[0] * upscale_factor // 8 * 8
target_height = lq_image.size[1] * upscale_factor // 8 * 8
resized_image = lq_image.resize((target_width, target_height))

# 3. Generate High-Resolution Image
output = pipeline(
    t_max=0.6667,
    t_min=0.0,
    tile_diffusion=True,              # Saves VRAM on large images
    tile_size=512,
    tile_stride=256,
    prompt="clean, high resolution, sharp details",
    negative_prompt="blurry, low quality, noise, artifacts",
    image=resized_image,
    num_inference_steps=10,
    guidance_scale=5.0,
    conditioning_scale=1.0,
    start_steps=19,
    start_point="lr",
    use_vae_encode_condition=True,
)

# 4. Save result
output.images[0].save("output_sr.png")

🎓 Citation & Original Work

This work is based on the paper "Improving the stability and efficiency of diffusion models for content consistent super-resolution". If you use these models, please cite the original authors:

@article{sun2024improving,
  title={Improving the stability and efficiency of diffusion models for content consistent super-resolution},
  author={Sun, Lingchen and Wu, Rongyuan and Liang, Jie and Zhang, Zhengqiang and Yong, Hongwei and Zhang, Lei},
  journal={arXiv preprint arXiv:2401.00877},
  year={2024}
}