CCSR: Improving the Stability and Efficiency of Diffusion Models for Content Consistent Super-Resolution

This repository contains CCSR-v2 (pruned), packaged as a standard, installable Python library (ccsr-pruned).

• Original Authors: Lingchen Sun, Rongyuan Wu, Jie Liang, Zhengqiang Zhang, Hongwei Yong, and Lei Zhang (The Hong Kong Polytechnic University & OPPO Research Institute).
• Original Paper: arXiv:2401.00877
• Original GitHub Repository: csslc/CCSR (CCSR-v2.0 branch)

🛠️ Installation

You can install this library using either pip or uv:

Option A: From PyPI (Recommended)

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

Option B: Directly from GitHub (Latest)

pip install git+https://github.com/AI-Wrappers/ccsr-v2-pruned.git
# or using uv
uv add git+https://github.com/AI-Wrappers/ccsr-v2-pruned.git

📦 Model Weights

All required model weights (including pre-trained Stable Diffusion 2.1 Base, custom VAE, and ControlNet) can be downloaded here: 👉 HuggingFace Hub: kharma1/ccsr_v2_repost

Make sure to download:

• stable-diffusion-2-1-base (standard SD 2.1 components).
• controlnet weights (CCSR-specific ControlNet model).
• vae weights (CCSR-specific fine-tuned Stage 2 VAE).

🚀 Quick Start / Usage

Here is a clean, simplified example of how to import the library and run super-resolution (SR) using the high-level CCSRUpscaler class, which automatically handles VAE tiling, image resizing, pipeline execution, and post-processing color fixes.

import torch
from PIL import Image
from ccsr import CCSRUpscaler

# 1. Setup model repository path
model_repo = "kharma1/ccsr_v2_repost"

# 2. Instantiate the upscaler (loads all subfolders directly from HF Hub)
upscaler = CCSRUpscaler(
    controlnet=(model_repo, "controlnet"),
    vae=(model_repo, "vae"),
    unet=(model_repo, "unet"),
    text_encoder=(model_repo, "text_encoder"),
    tokenizer=(model_repo, "tokenizer"),
    feature_extractor=(model_repo, "feature_extractor"),
    scheduler=(model_repo, "scheduler"),
    sample_method="ddpm",
    mixed_precision="fp16",
    tile_vae=True
)

# 4. Load low-quality (LQ) image and upscale
lq_image = Image.open("input_lq.png").convert("RGB")

sr_image = upscaler.upscale(
    image=lq_image,
    upscale=4,                        # Target upscale factor (e.g. 4x)
    num_inference_steps=6,            # Denoising steps
    t_max=0.6666,                     # Start noise level
    t_min=0.5,                        # Stop noise level (e.g. 0.5 for fast 1-step sampling)
    align_method="adain"              # Color fixing algorithm ('adain', 'wavelet', or 'nofix')
)

# 5. Save the result
sr_image.save("output_sr.png")

⚙️ CCSRUpscaler API Reference

Constructor Parameters (CCSRUpscaler(...))

The CCSRUpscaler class constructor sets up the accelerator and loads all model components.

Note: PathOrRepoParam is defined as Union[str, Tuple[str, Optional[str]], Tuple[str, Optional[str], Optional[str]]]. This permits passing:

• A simple local path string (e.g., /workspace/models)
• A repo tuple with subfolder (e.g., (model_repo, "vae"))
• A repo tuple with subfolder and a specific weights variant (e.g., (model_repo, "unet", "fp16"))

Upscaling Parameters (upscaler.upscale(...))

The upscale() method executes the super-resolution pipeline.

🛠️ Advanced Usage (Low-level Pipeline API)

If you prefer to directly interact with the underlying pipeline, you can import and call StableDiffusionControlNetCCSRPipeline directly:

from ccsr import StableDiffusionControlNetCCSRPipeline, ControlNetCCSRModel

# Load components
controlnet = ControlNetCCSRModel.from_pretrained("/workspace/models", subfolder="controlnet")
pipeline = StableDiffusionControlNetCCSRPipeline.from_pretrained("isometricneko/stable-diffusion-v2.1-clone", controlnet=controlnet)

# Run raw pipeline (requires manual resizing beforehand)
output = pipeline(
    t_max=0.6667,
    t_min=0.5,
    prompt="clean, high resolution",
    image=resized_image,
    num_inference_steps=6,
)

🎓 Citation

If you find CCSR helpful in your research or projects, please cite the original paper:

@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}
}