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.

Parameter	Type	Default	Description
controlnet	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load custom ControlNet.
vae	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load fine-tuned Stage 2 VAE.
unet	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load UNet model.
text_encoder	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load Text Encoder.
tokenizer	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load Tokenizer.
feature_extractor	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load Feature Extractor.
scheduler	PathOrRepoParam	Required	Path or (repo_id, subfolder) to load Scheduler.
sample_method	str	"ddpm"	Diffusion sampler. Choices: "ddpm", "unipcmultistep", "dpmmultistep".
mixed_precision	str	"fp16"	Mixed precision mode. Choices: "no", "fp16", "bf16".
tile_vae	bool	True	Enable tiling inside the custom VAE (saves VRAM).
vae_encoder_tile_size	int	1024	Tile size for VAE encoder.
vae_decoder_tile_size	int	224	Tile size for VAE decoder.
accelerator	Accelerator	None	Optional pre-configured HF Accelerator object.

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"))

Tip

Difference between variant and mixed_precision:

• variant (specified as the third element of the tuple, e.g., "fp16") controls download size and disk storage. It tells the loader to fetch specific files from the Hub (such as *.fp16.safetensors instead of *.safetensors). Useful for components like UNet and Text Encoder to save download time and disk space.
• mixed_precision (set in the constructor, e.g., "fp16") controls runtime GPU execution precision. It instructs the Accelerator to run calculations in half-precision and casts the loaded weights in-memory (e.g., converting loaded fp32 models to float16 on the GPU).

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

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

Parameter	Type	Default	Description
image	PIL.Image	Required	The low-resolution PIL Image.
prompt	str	"clean, ..."	Text prompt to guide details generation.
negative_prompt	str	"blurry, ..."	Negative prompt to avoid artifacts.
num_inference_steps	int	6	Total timeline steps configured for the scheduler.
guidance_scale	float	1.0	Classifier-Free Guidance (CFG). Set higher to guide more via prompt.
conditioning_scale	float	1.0	Strength of ControlNet alignment with input structure.
upscale	int	4	Target upscaling factor multiplier.
process_size	int	512	Minimum image dimension threshold (pre-rescales if too small).
t_max	float	0.6666	Start noise scale along the scheduler timeline.
t_min	float	0.5	Stop noise scale along the scheduler timeline.
start_steps	int	999	Timestep index threshold when noising input.
start_point	str	"lr"	Initialization method. "lr" (low-res latents) or "noise".
align_method	str	"adain"	Post-processing color profile correction. "adain", "wavelet", or "nofix".
tile_diffusion	bool	False	Enable sliding window tiling for UNet diffusion loops (saves VRAM).
tile_diffusion_size	int	512	Tile size for diffusion loop.
tile_diffusion_stride	int	256	Overlap stride for diffusion tiles.
use_vae_encode_condition	bool	True	Encodes low-res condition via VAE. Must be True for optimal CCSR v2 quality.
seed	int	None	Optional seed for reproducibility.
sample_times	int	1	Number of samples to generate. Returns a list if > 1.

---

🛠️ 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}
}