Metadata-Version: 2.4
Name: ccsr-pruned
Version: 2.1.0
Summary: Continuous Latent Diffusion for Image Super-Resolution (CCSR) packaged as an installable library
Author: AI-Wrappers
License: Apache-2.0
License-File: LICENSE
Keywords: ccsr,controlnet,diffusers,diffusion,image-upscaling,super-resolution
Requires-Python: >=3.13
Requires-Dist: accelerate>=1.14.0
Requires-Dist: diffusers>=0.39.0
Requires-Dist: einops>=0.8.2
Requires-Dist: numpy>=2.5.1
Requires-Dist: pillow>=12.3.0
Requires-Dist: safetensors>=0.8.0
Requires-Dist: torch>=2.13.0
Requires-Dist: torchvision>=0.28.0
Requires-Dist: transformers>=5.13.0
Description-Content-Type: text/markdown

# 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](https://scholar.google.com/citations?hl=zh-CN&tzom=-480&user=ZCDjTn8AAAAJ), [Rongyuan Wu](https://scholar.google.com/citations?user=A-U8zE8AAAAJ&hl=zh-CN), [Jie Liang](https://scholar.google.com.sg/citations?user=REWxLZsAAAAJ&hl), [Zhengqiang Zhang](https://scholar.google.com/citations?hl=zh-CN&user=UX26wSMAAAAJ&view_op=list_works&sortby=pubdate), [Hongwei Yong](https://scholar.google.com.hk/citations?user=Xii74qQAAAAJ&hl=zh-CN), and [Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang) (The Hong Kong Polytechnic University & OPPO Research Institute).
* **Original Paper:** [arXiv:2401.00877](https://arxiv.org/pdf/2401.00877)
* **Original GitHub Repository:** [csslc/CCSR (CCSR-v2.0 branch)](https://github.com/csslc/CCSR/tree/CCSR-v2.0)

---

## 🛠️ Installation

You can install this library using either `pip` or `uv`:

### Option A: From PyPI (Recommended)
```bash
pip install ccsr-pruned
# or using uv
uv add ccsr-pruned
```

### Option B: Directly from GitHub (Latest)
```bash
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](https://huggingface.co/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) on an image using tiled processing to save VRAM:

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

# 1. Load the model components (replace paths with your download locations)
model_path = "preset/models/stable-diffusion-2-1-base"
controlnet_path = "preset/models"  # Path to the directory containing controlnet/
vae_path = "preset/models"         # Path to the directory containing vae/

# Load ControlNet & VAE
controlnet = ControlNetCCSRModel.from_pretrained(controlnet_path, subfolder="controlnet", torch_dtype=torch.float16)
vae = AutoencoderKL.from_pretrained(vae_path, subfolder="vae", torch_dtype=torch.float16)

# Initialize the pipeline
pipeline = StableDiffusionControlNetCCSRPipeline.from_pretrained(
    model_path,
    controlnet=controlnet,
    vae=vae,
    torch_dtype=torch.float16,
).to("cuda")

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

# The pipeline expects the input image to be resized to the target SR resolution (and divisible by 8)
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. Run Super-Resolution
# Note: CCSR-v2 supports flexible diffusion steps (e.g. 10 steps, 2 steps, or 1 step)
output = pipeline(
    t_max=0.6667,                     # Start noise level
    t_min=0.0,                        # End noise level
    tile_diffusion=True,              # Set to True to enable tiled processing (saves VRAM)
    tile_size=512,                    # Resolution tile size (e.g., 512)
    tile_stride=256,                  # Overlap stride (e.g., 256)
    prompt="clean, high resolution, sharp details",
    negative_prompt="blurry, low quality, noise, artifacts",
    image=resized_image,
    num_inference_steps=10,           # Total inference steps
    guidance_scale=5.0,               # CFG scale
    conditioning_scale=1.0,           # ControlNet strength
    start_steps=19,                   # Starting timestep index for Stage 2
    start_point="lr",                 # 'lr' (low-res) or 'vae'
    use_vae_encode_condition=True,    # Use VAE to encode condition
)

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

---

## ⚙️ Configuration Parameters

The `pipeline(...)` call accepts several parameters to tune the output and resource usage:

| Parameter | Type | Default | Description |
| :--- | :--- | :--- | :--- |
| `t_max` | `float` | - | Maximum noise timestep scale to start sampling (e.g., `0.6667` for 10-step). |
| `t_min` | `float` | - | Minimum noise timestep scale to end sampling (typically `0.0`). |
| `tile_diffusion` | `bool` | - | Enables tiled diffusion. Essential for upscaling large images without running out of GPU memory. |
| `tile_size` | `float` | - | Width/height of the processing tile (e.g., `512` or `256`). |
| `tile_stride` | `float` | - | Step size between tiles (typically half of `tile_size`). |
| `image` | `PIL.Image` | - | Low-quality input image **already resized to the target upscaled resolution** (divisible by 8). |
| `guidance_scale` | `float` | `7.5` | Classifier-Free Guidance (CFG) scale. Higher values increase prompt fidelity but can introduce artifacts. |
| `conditioning_scale` | `float` | `1.0` | Strength of the ControlNet model. |
| `start_steps` | `int` | `999` | Starting step threshold. |
| `start_point` | `str` | `'noise'` | Can be `'noise'`, `'lr'`, or `'vae'` depending on your initialization method. |
| `use_vae_encode_condition`| `bool` | `False` | Encodes the low-resolution condition via VAE. |

---

## 🎓 Citation

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

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