Metadata-Version: 2.4
Name: ccsr-pruned
Version: 2.2.1
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: 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
Provides-Extra: test
Requires-Dist: tqdm>=4.66.0; extra == 'test'
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) using the high-level `CCSRUpscaler` class, which automatically handles VAE tiling, image resizing, pipeline execution, and post-processing color fixes.

```python
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:

```python
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:

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