### Install VOSR Dependencies

Source: https://github.com/cswry/vosr/blob/main/README.md

Commands to clone the repository and set up the required Conda environment.

```bash
## clone this repository
git clone https://github.com/cswry/VOSR.git
cd VOSR

# create an environment with python >= 3.8
conda create -n vosr python=3.8
conda activate vosr
pip install -r requirements.txt
```

--------------------------------

### Train VOSR-1.4B Multi-step Model

Source: https://context7.com/cswry/vosr/llms.txt

Start training for the VOSR-1.4B multi-step model, utilizing 8 GPUs for distributed training. The training configuration is specified in a YAML file.

```bash
torchrun --nproc_per_node=8 train_vosr.py \
    --config configs/train_yml/multi_step/VOSR_1.4B.yml
```

--------------------------------

### Applying Color Alignment in Inference

Source: https://context7.com/cswry/vosr/llms.txt

Example of how to apply either AdaIN or Wavelet color correction during the inference process. Assumes sr_img and input_img are Pillow Image objects and align_method is a string ('adain' or 'wavelet').

```python
# Example usage in inference
sr_img = transforms.ToPILImage()(sr_tensor[0].cpu() * 0.5 + 0.5)

# Apply color alignment
if align_method == 'adain':
    sr_img = adain_color_fix(sr_img, input_img)
elif align_method == 'wavelet':
    sr_img = wavelet_color_fix(sr_img, input_img)

sr_img.save('output.png')
```

--------------------------------

### Train VOSR-1.4B One-step Distilled Model

Source: https://context7.com/cswry/vosr/llms.txt

Train the VOSR-1.4B one-step distilled model. This process leverages a pre-trained multi-step model for distillation and requires distributed training setup.

```bash
torchrun --nproc_per_node=8 train_vosr_distill.py \
    --config configs/train_yml/one_step/VOSR_1.4B.yml
```

--------------------------------

### Initialize VOSR and Perform Flow Matching

Source: https://context7.com/cswry/vosr/llms.txt

Configures the VOSR environment and demonstrates multi-step and one-step sampling procedures.

```python
vosr = VOSR(
    time_dist=['lognorm', -0.4, 1.0],  # Time distribution: lognorm or uniform
    cfg_ratio=0.10,                     # CFG training ratio (10% samples use weak condition)
    cfg_scale=2.0,                      # CFG inference scale
    interp_type='lin',                  # Interpolation type: 'lin' or 'sph'
    accelerator=accelerator,            # Hugging Face Accelerator
    t_start=0.0,                        # CFG start time
    t_end=1.0,                          # CFG end time
    args=args                           # Config namespace
)

# Compute multi-step flow matching loss
# lq: low-quality latent [B, C, H, W], hq: high-quality latent [B, C, H, W]
# z: DINOv2 features list
loss, loss_backward = vosr.loss_fm(model, lq, hq, z)

# Sample with multi-step flow matching (25 steps default)
with torch.no_grad():
    sr_latent = vosr.sample_multistep_fm(
        model=model,
        lq=lq_latent,
        venc_fea=dinov2_features,
        n_steps=25,
        schedule='linear'  # 'linear' or 'cosine'
    )

# One-step sampling (for distilled models)
with torch.no_grad():
    sr_latent = vosr.sample_onestep(
        model=model,
        lq=lq_latent,
        venc_fea=dinov2_features,
        n_steps=1,
        schedule='linear'
    )
```

--------------------------------

### Execute Multi-step Training

Source: https://github.com/cswry/vosr/blob/main/README.md

Run multi-step training for VOSR models using torchrun.

```bash
# VOSR-0.5B
torchrun --nproc_per_node=8 train_vosr.py --config configs/train_yml/multi_step/VOSR_0.5B.yml

# VOSR-1.4B
torchrun --nproc_per_node=8 train_vosr.py --config configs/train_yml/multi_step/VOSR_1.4B.yml
```

--------------------------------

### Configure Training Dataset

Source: https://github.com/cswry/vosr/blob/main/README.md

Define folder paths and sampling weights for training data in a text configuration file.

```text
/path/to/dataset_A, 2
/path/to/dataset_B, 1
/path/to/dataset_C, 1
```

--------------------------------

### Initialize Data Loaders

Source: https://context7.com/cswry/vosr/llms.txt

Python implementation for parsing dataset configuration files and initializing PyTorch DataLoaders for both txt and webdataset modes.

```python
from dataloaders.realsr_dataset import TxtPairDataset, build_webdataset_pipeline
from torch.utils.data import DataLoader

# Load dataset configuration
def load_dataset_config(path):
    txt_paths = []
    probs = []
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = [x.strip() for x in line.split(",")]
            dataset_path = parts[0]
            repeat = int(parts[1]) if len(parts) == 2 else 1
            txt_paths.append(dataset_path)
            probs.append(repeat)
    return txt_paths, probs

# Create dataloader for txt mode
train_dataset = TxtPairDataset(split='train', args=args)
train_dataloader = DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True,
    prefetch_factor=4
)

# Create dataloader for webdataset mode
train_dataset = build_webdataset_pipeline(args, split='train')
train_dataloader = DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=False,  # WebDataset handles shuffling
    num_workers=8,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True,
    prefetch_factor=4
)
```

--------------------------------

### Execute One-step Distillation

Source: https://github.com/cswry/vosr/blob/main/README.md

Run one-step distillation training requiring a pre-trained multi-step teacher checkpoint.

```bash
# VOSR-0.5B one-step
torchrun --nproc_per_node=8 train_vosr_distill.py --config configs/train_yml/one_step/VOSR_0.5B.yml

# VOSR-1.4B one-step
torchrun --nproc_per_node=8 train_vosr_distill.py --config configs/train_yml/one_step/VOSR_1.4B.yml
```

--------------------------------

### Configure LightningDiT Model

Source: https://context7.com/cswry/vosr/llms.txt

Initializes the LightningDiT transformer backbone for multi-step or distilled training and demonstrates forward passes.

```python
from models.lightningdit import LightningDiT

# Initialize LightningDiT for multi-step training (no auxiliary time condition)
model = LightningDiT(
    input_size=64,           # Latent size (512 / 8 = 64)
    patch_size=2,            # Patch size for transformer
    in_channels=8,           # 2 * latent_channels (lq + zt concatenated)
    out_channels=4,          # Latent channels
    hidden_size=1024,        # Transformer hidden dimension
    depth=28,                # Number of transformer blocks
    num_heads=16,            # Attention heads
    mlp_ratio=4,             # MLP expansion ratio
    z_dims=768,              # DINOv2 feature dimension
    encdim_ratio=3,          # Cross-attention MLP ratio
    auxiliary_time_cond=False,  # False for multi-step, True for distilled
    use_qknorm=True,         # Query-Key normalization
    use_swiglu=True,         # SwiGLU activation in MLP
    use_rope=True,           # Rotary Position Embedding
    use_rmsnorm=True,        # RMSNorm instead of LayerNorm
    wo_shift=False,          # Disable shift in AdaLN
    num_fused_layers=1       # Number of DINOv2 layers to fuse
)

# Initialize for one-step distilled model (with auxiliary time condition)
distilled_model = LightningDiT(
    input_size=64,
    patch_size=2,
    in_channels=8,
    out_channels=4,
    hidden_size=1024,
    depth=28,
    num_heads=16,
    mlp_ratio=4,
    z_dims=768,
    encdim_ratio=3,
    auxiliary_time_cond=True,  # Enable for distilled models
    use_qknorm=True,
    use_swiglu=True,
    use_rope=True,
    use_rmsnorm=True,
    wo_shift=False,
    num_fused_layers=1
)

# Forward pass (multi-step model)
# inp: [B, in_channels, H, W] - concatenated [lq_latent, z_t]
# t: [B] - timestep
# z: list of DINOv2 features
output = model(inp, t, z=dinov2_features)

# Forward pass (distilled model with auxiliary time)
# r: [B] - auxiliary timestep (shortcut endpoint)
output = distilled_model(inp, t=t_current, r=t_next, z=dinov2_features)

# Flexible forward for variable input sizes
model.forward = model.forward_flexible
output = model(inp, t, z=dinov2_features)
```

--------------------------------

### Run Multi-step Inference

Source: https://github.com/cswry/vosr/blob/main/README.md

Perform inference using multi-step models with configurable sampling steps and conditioning scales.

```bash
# Inputs under preset/datasets/inp_data follow RealSR-style evaluation; use --cfg_scale -0.5 (see benchmark presets above).

# VOSR-0.5B multi-step (25 steps)
python inference_vosr.py \
    -c preset/ckpts/VOSR_0.5B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

# VOSR-1.4B multi-step (25 steps)
python inference_vosr.py \
    -c preset/ckpts/VOSR_1.4B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4
```

--------------------------------

### Basic 4x Upscaling with VOSR-0.5B Multi-step

Source: https://context7.com/cswry/vosr/llms.txt

Perform basic 4x super-resolution using the VOSR-0.5B multi-step model. Ensure the checkpoint and input data paths are correctly specified.

```bash
python inference_vosr.py \
    -c preset/ckpts/VOSR_0.5B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4
```

--------------------------------

### Train VOSR-0.5B Multi-step Model

Source: https://context7.com/cswry/vosr/llms.txt

Initiate training for the VOSR-0.5B multi-step model using distributed data parallel on 8 GPUs. Configuration is managed via a YAML file.

```bash
torchrun --nproc_per_node=8 train_vosr.py \
    --config configs/train_yml/multi_step/VOSR_0.5B.yml
```

--------------------------------

### VOSR-1.4B Multi-step with Custom CFG for ScreenSR

Source: https://context7.com/cswry/vosr/llms.txt

Run VOSR-1.4B multi-step inference with a custom classifier-free guidance (CFG) scale, suitable for benchmarks like ScreenSR. Adjust `cfg_scale` and `weak_cond_strength_aelq` for specific quality-generation trade-offs.

```bash
python inference_vosr.py \
    -c preset/ckpts/VOSR_1.4B_ms \
    -i /path/to/ScreenSR/test_images \
    -o preset/results \
    -u 4 \
    --cfg_scale 0.5 \
    --weak_cond_strength_aelq 0.1
```

--------------------------------

### Implement Distillation Loss Functions

Source: https://context7.com/cswry/vosr/llms.txt

Trains a student model to match teacher outputs using shortcut consistency or RCGM methods.

```python
from vosr import VOSR

# Initialize VOSR for distillation training
vosr = VOSR(
    u_weight=1.0,                       # Weight for consistency loss
    time_dist=['uniform', 0, 1.0],
    cfg_ratio=0.1,
    cfg_scale=0.5,
    interp_type='lin',
    accelerator=accelerator,
    t_start=0,
    t_end=1,
    args=args
)

# Shortcut consistency distillation loss
loss, loss_backward = vosr.loss_fm_distill_shortcut_improved(
    model=student_model,
    lq=lq_latent,
    hq=hq_latent,
    z=dinov2_features,
    model_tea=teacher_model
)

# RCGM consistency distillation loss
loss, loss_backward = vosr.loss_fm_distill_rcgm_improved(
    model=student_model,
    lq=lq_latent,
    hq=hq_latent,
    z=dinov2_features,
    model_tea=teacher_model
)

# Backward and optimize
accelerator.backward(loss_backward)
optimizer.step()
```

--------------------------------

### Resume Multi-step Training from Checkpoint

Source: https://context7.com/cswry/vosr/llms.txt

Continue a previously interrupted multi-step training process for the VOSR-0.5B model. The training script automatically detects the latest checkpoint in the output directory.

```bash
torchrun --nproc_per_node=8 train_vosr.py \
    --config configs/train_yml/multi_step/VOSR_0.5B.yml
# (Auto-discovers latest checkpoint in output_dir/checkpoints/)
```

--------------------------------

### Configure VOSR Training Parameters

Source: https://context7.com/cswry/vosr/llms.txt

YAML configuration file defining model architecture, hyperparameters, and training settings.

```yaml
# configs/train_yml/multi_step/VOSR_0.5B.yml

# Dataset configuration
train_dataset_config: configs/train_txt/train_dataset_txt.txt
dataset_type: 'webdataset'  # 'txt' or 'webdataset'
shuffle_buffer: 2000

# Test dataset for evaluation during training
test_lq_dir: preset/datasets/RealSR/test_LR
test_gt_dir: preset/datasets/RealSR/test_HR
test_upscale: 4

# Model architecture
resolution: 512
patch_size: 2
dim: 1024           # hidden_size
depth: 28           # number of transformer blocks
num_heads: 16
mlp_ratio: 4
use_qknorm: True
use_swiglu: True
use_rope: True
use_rmsnorm: True
wo_shift: False

# VAE configuration
ae_type: sd2        # 'sd2' for 0.5B, 'qwen' for 1.4B
ae_path: preset/ckpts/stable-diffusion-2-1-base
lwdecoder_path: preset/ckpts/sd21_lwdecoder.pth

# DINOv2 encoder configuration
dinov2_size: 448
enc_type: dinov2b   # 'dinov2b', 'dinov2l', or 'dinov2g'
enc_dim: 768
layer_dinov2b_list: [8]
encdim_ratio: 3

# Training hyperparameters
train_batch_size: 8
gradient_accumulation_steps: 1
learning_rate: 5.0e-5
max_train_steps: 400001
lr_scheduler: constant_with_warmup
lr_warmup_steps: 0
adam_beta1: 0.9
adam_beta2: 0.95
adam_weight_decay: 1.0e-2
max_grad_norm: 1.0
mixed_precision: bf16
use_ema: True
ema_decay: 0.9999
seed: 42

# Flow matching configuration
time_dist: ['uniform', 0, 1.0]  # or ['lognorm', -0.4, 1.0]
cfg_ratio: 0.1
cfg_scale: -0.5
weak_cond_strength_aelq_list: [0.05, 0.25]
cond_strength_aelq_list: [5, 1]
interp_type: lin
t_start: 0
t_end: 1

# Inference during training
infer_steps: 25
checkpointing_steps: 10000
inference_steps: 10000

# Checkpoint loading
resume_ckpt: ~      # Auto-discover or specify path
pretrained_ckpt: ~  # Initialize from pretrained weights

# Logging
output_dir: exp_vosr/
logging_dir: logs
report_to: none     # 'wandb' or 'none'
tracker_project_name: vosr
```

--------------------------------

### Define Dataset Paths and Weights

Source: https://context7.com/cswry/vosr/llms.txt

Text file format for specifying dataset locations and their relative sampling weights.

```text
# configs/train_txt/train_dataset_txt.txt
# Format: <folder_path>, <sampling_weight>
# Higher weight = more frequent sampling

/path/to/dataset_A, 2
/path/to/dataset_B, 1
/path/to/dataset_C, 1
```

--------------------------------

### One-step Tiled Inference for Large Images

Source: https://context7.com/cswry/vosr/llms.txt

Apply one-step super-resolution to large images using tiled processing to manage memory. Configure `tile_size` and `tile_overlap` for optimal performance.

```bash
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_0.5B_os \
    -i /path/to/4k_image.png \
    -o preset/results \
    -u 4 \
    --tile_size 512 \
    --tile_overlap 32
```

--------------------------------

### Multi-step Inference with Tiled Processing for Large Images

Source: https://context7.com/cswry/vosr/llms.txt

Process large images using tiled inference to manage memory usage. Specify `tile_size` and `tile_overlap` for effective tiling. The `infer_steps` parameter is set to 25 for multi-step inference.

```bash
python inference_vosr.py \
    -c preset/ckpts/VOSR_0.5B_ms \
    -i /path/to/large_image.png \
    -o preset/results \
    -u 4 \
    --tile_size 512 \
    --tile_overlap 32 \
    --infer_steps 25
```

--------------------------------

### Model Weights Directory Structure

Source: https://github.com/cswry/vosr/blob/main/README.md

The expected file layout for pretrained checkpoints and VAE assets within the preset/ckpts/ directory.

```text
preset/ckpts/
|-- Qwen-Image-vae-2d/          # Qwen-Image VAE (2D, for 1.4B models)
|-- stable-diffusion-2-1-base/  # SD2.1 VAE (for 0.5B models)
|-- sd21_lwdecoder.pth          # Lightweight decoder for SD2.1 VAE
|-- torch_cache/                # DINOv2 pretrained weights
|-- VOSR_0.5B_ms/               # 0.5B multi-step model
|-- VOSR_0.5B_os/               # 0.5B one-step (distilled) model
|-- VOSR_1.4B_ms/               # 1.4B multi-step model
`-- VOSR_1.4B_os/               # 1.4B one-step (distilled) model
```

--------------------------------

### Run One-step Inference

Source: https://github.com/cswry/vosr/blob/main/README.md

Perform inference using distilled one-step models.

```bash
# VOSR-0.5B one-step
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_0.5B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

# VOSR-1.4B one-step
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_1.4B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4
```

--------------------------------

### VAE Encoding and Decoding Pipelines

Source: https://context7.com/cswry/vosr/llms.txt

Utilities for loading VAE backends and performing latent space transformations. Supports both Stable Diffusion 2.1 and Qwen-Image architectures.

```python
import torch
from diffusers import AutoencoderKL
from models.qwenimage_vae2d import AutoencoderKLQwenImage2D
from models.light_decoder import LightDecoder

# Load SD2.1 VAE for VOSR-0.5B
def load_sd2_vae(ae_path, decoder_path, device):
    vae = AutoencoderKL.from_pretrained(ae_path, subfolder="vae")
    vae.to(device).eval()

    # Load lightweight decoder for faster decoding
    ckpt = torch.load(decoder_path, map_location="cpu")
    light_decoder = LightDecoder(
        in_channels=ckpt["config"]["in_channels"],
        out_channels=ckpt["config"]["out_channels"],
        block_out_channels=tuple(ckpt["config"]["block_out_channels"]),
        layers_per_block=ckpt["config"]["layers_per_block"],
    )
    light_decoder.load_state_dict(ckpt["model_state_dict"])
    light_decoder.to(device).eval()

    return vae, light_decoder

# Load Qwen-Image VAE for VOSR-1.4B
def load_qwen_vae(ae_path, device):
    vae = AutoencoderKLQwenImage2D.from_pretrained(ae_path)
    vae.to(device).eval()
    return vae

# Encode image to latent (SD2.1)
def encode_sd2(vae, x):
    """x: [B, 3, H, W] in range [-1, 1]"""
    latent = vae.encode(x).latent_dist.sample()
    return latent * vae.config.scaling_factor

# Encode image to latent (Qwen-Image)
def encode_qwen(vae, x, device):
    """x: [B, 3, H, W] in range [-1, 1]"""
    latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1).to(device)
    latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, -1, 1, 1).to(device)
    latent = vae.encode(x).latent_dist.sample()
    return (latent - latents_mean) * latents_std

# Decode latent to image (SD2.1 with light decoder)
def decode_sd2(vae, light_decoder, latent):
    sr_u = latent / vae.config.scaling_factor
    return light_decoder(sr_u).clamp(-1, 1)

# Decode latent to image (Qwen-Image)
def decode_qwen(vae, latent, latents_mean, latents_std):
    latent = latent / latents_std + latents_mean
    return vae.decode(latent, return_dict=False)[0].clamp(-1, 1)

# Example: Full encoding/decoding pipeline
device = 'cuda'
vae, light_decoder = load_sd2_vae(
    'preset/ckpts/stable-diffusion-2-1-base',
    'preset/ckpts/sd21_lwdecoder.pth',
    device
)

# Encode LQ image
lq_image = torch.randn(1, 3, 512, 512).to(device)  # [-1, 1]
lq_latent = encode_sd2(vae, lq_image)

# After SR processing...
sr_latent = lq_latent  # placeholder

# Decode to image
sr_image = decode_sd2(vae, light_decoder, sr_latent)  # [-1, 1]
sr_image_01 = (sr_image + 1) / 2  # Convert to [0, 1] for saving
```

--------------------------------

### RealSR Benchmark Preset Inference

Source: https://context7.com/cswry/vosr/llms.txt

Execute inference using parameters optimized for the RealSR benchmark, which involves a negative `cfg_scale` value. This setting influences the guidance strength during generation.

```bash
python inference_vosr.py \
    -c preset/ckpts/VOSR_0.5B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4 \
    --cfg_scale -0.5
```

--------------------------------

### Initialize VOSR Class

Source: https://context7.com/cswry/vosr/llms.txt

Import and initialize the VOSR class in Python. This class encapsulates the core logic for training losses, sampling methods, and generation, including classifier-free guidance and condition interpolation.

```python
import torch
from vosr import VOSR
```

--------------------------------

### One-step Inference with Custom Seed

Source: https://context7.com/cswry/vosr/llms.txt

Generate super-resolved images with a specific seed for reproducible results in one-step inference. The `--seed` parameter ensures consistent outputs across multiple runs.

```bash
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_0.5B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4 \
    --seed 123
```

--------------------------------

### One-step Inference with VOSR-1.4B Distilled Model

Source: https://context7.com/cswry/vosr/llms.txt

Utilize the VOSR-1.4B one-step distilled model for efficient single-step super-resolution. This provides a balance between speed and quality for larger model variants.

```bash
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_1.4B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4
```

--------------------------------

### Train VOSR-0.5B One-step Distilled Model with RCGM

Source: https://context7.com/cswry/vosr/llms.txt

Train a VOSR-0.5B one-step distilled model using the RCGM consistency distillation method. This requires a pre-trained multi-step teacher model.

```bash
torchrun --nproc_per_node=8 train_vosr_distill.py \
    --config configs/train_yml/one_step/VOSR_0.5B.yml
```

--------------------------------

### Load and Extract DINOv2 Features

Source: https://context7.com/cswry/vosr/llms.txt

Functions to initialize the DINOv2 encoder and extract intermediate layer features from low-quality image tensors. The encoder parameters are frozen during extraction.

```python
import torch
from torchvision.transforms import Normalize

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)

def load_dinov2(enc_type, device):
    """Load DINOv2 encoder for visual feature extraction."""
    if enc_type == 'dinov2b':
        encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
    elif enc_type == 'dinov2l':
        encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
    elif enc_type == 'dinov2g':
        encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14')

    encoder.head = torch.nn.Identity()
    encoder = encoder.to(device).eval()

    for p in encoder.parameters():
        p.requires_grad_(False)

    return encoder

def preprocess_for_dinov2(x, target_size=448):
    """Preprocess image tensor for DINOv2 input."""
    x = x / 255.0
    x = torch.nn.functional.interpolate(x, target_size, mode='bicubic').clip(0., 1.)
    x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
    return x

def get_venc_features(venc, lq_tensor, dinov2_size=448, layer_list=[8]):
    """Extract DINOv2 features from low-quality image."""
    with torch.no_grad():
        # Convert from [-1, 1] to [0, 255]
        raw_image = (0.5 * lq_tensor + 0.5) * 255
        raw_image = preprocess_for_dinov2(raw_image, dinov2_size)

        # Get intermediate layer features
        features = {}
        x = venc.prepare_tokens_with_masks(raw_image, None)
        for i, blk in enumerate(venc.blocks):
            x = blk(x)
            features[f'layer_{i}'] = x[:, 1:]  # Remove CLS token

        x_norm = venc.norm(x)[:, 1:]
        z = [features[f'layer_{i}'] for i in layer_list]
        z[-1] = x_norm  # Use normalized output for last layer

    return z

# Example usage
device = 'cuda'
venc = load_dinov2('dinov2b', device)

# lq_tensor: [B, 3, H, W] in range [-1, 1]
lq_tensor = torch.randn(1, 3, 512, 512).to(device)
dinov2_features = get_venc_features(venc, lq_tensor, dinov2_size=448, layer_list=[8])
```

--------------------------------

### One-step Inference with VOSR-0.5B Distilled Model

Source: https://context7.com/cswry/vosr/llms.txt

Perform fast super-resolution using the VOSR-0.5B one-step distilled model. This mode is suitable for applications requiring real-time processing.

```bash
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_0.5B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4
```

--------------------------------

### Wavelet Color Correction Function

Source: https://context7.com/cswry/vosr/llms.txt

Performs color correction by preserving high-frequency details from the target image while matching low-frequency color from the source image using a wavelet-based approach. Requires OpenCV (cv2) and Pillow (PIL).

```python
import torch
import numpy as np
from PIL import Image

def wavelet_color_fix(target, source):
    """
    Wavelet-based color correction.
    Preserves high-frequency details from target while matching
    low-frequency color from source.
    """
    import cv2

    target_np = np.array(target).astype(np.float32) / 255.0
    source_np = np.array(source.resize(target.size, Image.LANCZOS)).astype(np.float32) / 255.0

    sigma = 5
    source_low = cv2.GaussianBlur(source_np, (0, 0), sigma)
    target_low = cv2.GaussianBlur(target_np, (0, 0), sigma)
    target_high = target_np - target_low

    result = np.clip(source_low + target_high, 0, 1) * 255.0
    return Image.fromarray(result.astype(np.uint8))
```

--------------------------------

### AdaIN Color Correction Function

Source: https://context7.com/cswry/vosr/llms.txt

Matches the mean and standard deviation of the target image to the source image using Adaptive Instance Normalization. Requires torchvision.transforms.functional.

```python
import torch
import numpy as np
from PIL import Image

def adain_color_fix(target, source):
    """
    AdaIN-based color correction.
    Matches mean and std of target to source.
    """
    from torchvision.transforms.functional import to_tensor, to_pil_image

    target_tensor = to_tensor(target).unsqueeze(0)
    source_tensor = to_tensor(source).unsqueeze(0)
    eps = 1e-5

    target_mean = torch.mean(target_tensor, dim=[2, 3], keepdim=True)
    target_std = torch.std(target_tensor, dim=[2, 3], keepdim=True) + eps
    source_mean = torch.mean(source_tensor, dim=[2, 3], keepdim=True)
    source_std = torch.std(source_tensor, dim=[2, 3], keepdim=True) + eps

    target_tensor = (target_tensor - target_mean) / target_std * source_std + source_mean
    return to_pil_image(torch.clamp(target_tensor[0], 0, 1))
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.