How 6ungying/piff-master02 works
1. Overall purpose
FloodDiff is a diffusion-model repository for image restoration / reconstruction experiments. The codebase is organized around generating clean images from degraded inputs, with multiple model families and degradation types.
2. Top-level execution flow
The main runtime entry points are train.py for model fitting, sample.py for generation/inference, compute_metrices.py for evaluation, and paper_vis.py for visualization outputs. scripts/train.sh and scripts/download_checkpoint.sh provide common automation around training and checkpoint retrieval.
3. Configuration boundary
Experiment behavior is driven by YAML configs in configs/ such as Template-BBDM.yaml, Template-LBBDM-f4.yaml, Template-LBBDM-f8.yaml, Template-LBBDM-f16.yaml, and flooding.yaml. These appear to define model variants, latent bottleneck settings, and dataset/degradation choices. checkpoint.json and the model storage docs indicate that pretrained weights are external artifacts the runtime expects to load.
4. Data layer
data.py and dataset/ define dataset loading, with an ImageNet-centric path visible in dataset/imagenet.py and validation split metadata files under dataset/. This layer is the source of clean images before corruption/degradation is applied.
5. Degradation / corruption subsystem
The corruption/ package encapsulates forward-image degradation models: blur, inpainting, JPEG, super-resolution, and mixtures. corruption/base.py defines the shared interface, while the specialized modules implement the actual transform. This is a key architectural boundary because the model learns to invert these degradations rather than modeling raw images only.
6. Main diffusion backbone
guided_diffusion/ is the core pixel-space diffusion implementation. It contains the UNet architecture, Gaussian diffusion process, attention blocks, losses, respacing utilities, FP16 helpers, and logging support. This package is the primary model/math engine for standard diffusion variants in the repo.
7. Alternative / extended model stack
i2sb/ contains a larger second modeling stack with its own diffusion modules, embeddings, network wrappers, physics-related loss, runner, checkpoint utilities, and utility code. The presence of i2sb/runner.py and i2sb/diffusion.py suggests a higher-level training/inference framework layered on top of lower-level diffusion blocks.
8. Latent-space branch via VQGAN
i2sb/VQGAN/ provides the latent compression side of the system: VQGAN model code, quantization, discriminator, perceptual loss, and a vendored taming/ implementation with datasets and transformer components. The Template-LBBDM-f4/f8/f16 configs strongly suggest this branch supports latent-bottleneck diffusion at different compression factors.
9. Evaluation and metrics
evaluation/ contains FID-related utilities and a ResNet-based evaluator. Combined with compute_metrices.py, this forms the scoring path for generated outputs, separating evaluation concerns from training and sampling.
10. Training support and infrastructure
distributed_util.py, guided_diffusion/fp16_util.py, guided_diffusion/logger.py, and top-level logger.py handle multi-device execution, mixed precision, and experiment logging. GPU.py, loss.py, loss_single.py, and mean_std.py look like supporting utilities for compute selection, custom losses, and dataset/statistics handling.
11. Architectural data flow to diagram
A typical flow is: clean dataset → optional corruption/degradation module → diffusion model (guided_diffusion or i2sb, possibly via VQGAN latent encoding) → sampled/restored image outputs → metric computation in evaluation/. Pretrained checkpoints sit beside this flow as external inputs consumed by training or inference.
12. Repository boundary summary
The repository is split into three major concerns: data/degradation, diffusion modeling, and evaluation/visualization. The pixel-space diffusion path (guided_diffusion) and the latent-space path (i2sb + VQGAN) are the most important model boundaries to show in an architecture diagram.
Main entry points: train.py train.py · runner.py Runner · sample.py sample.py · data.py data.py
Components 22
Entrypoints
Training, sampling, evaluation, and visualization scripts.
- train.py · Main model fitting script.
train.py - sample.py · Generation and restoration sampling script.
sample.py - compute_metrices.py · Runs quantitative evaluation over outputs.
compute_metrices.py - paper_vis.py · Produces figures and visualization artifacts.
paper_vis.py - configs · YAML templates for model variants, latent sizes, and degradation setups.
configs
Data & Degradation
Dataset loading and image corruption pipelines.
- data.py · Top-level dataset and dataloader orchestration.
data.py - ImageNet data · ImageNet-centric dataset implementation and validation split metadata.
dataset/imagenet.py - Corruption base · Shared interface for forward image corruption models.
corruption/base.py - Corruption ops · Blur, inpainting, JPEG, super-resolution, and mixtures.
corruption
Pixel Diffusion
Core guided diffusion stack for image restoration.
- UNet · Main neural network architecture for pixel-space diffusion.
guided_diffusion/unet.py - Diffusion core · Gaussian diffusion, respacing, and sampling math.
guided_diffusion/gaussian_diffusion.py - Guided losses · Loss functions and training objectives for diffusion models.
guided_diffusion/losses.py - Guided utils · Attention, fp16, logging, script helpers, and runtime utilities.
guided_diffusion
Latent Diffusion
I2SB and VQGAN-based latent model stack.
- Runner · Higher-level training and inference loop for the i2sb stack.
i2sb/runner.py - i2sb diffusion · Alternative diffusion modules and process definitions.
i2sb/diffusion.py - Network · Model wiring, embeddings, and network assembly.
i2sb/network.py - VQGAN · Latent compression and reconstruction model used for bottleneck diffusion.
i2sb/VQGAN/vqgan.py - Taming · Vendored VQGAN datasets, modules, and transformer utilities.
i2sb/VQGAN/taming
Evaluation
Metrics and scoring for generated outputs.
- Evaluation pkg · FID utilities and ResNet-based evaluator.
evaluation
Infrastructure
Logging, distributed execution, and utility support.
- Logging · Experiment logging and status reporting.
logger.py - Distributed · Multi-device execution helpers.
distributed_util.py - Checkpoints · External pretrained weights and checkpoint references.
checkpoint.json
Connections 24
- train.py → configs: reads · Loads experiment configuration.
- train.py → data.py: uses · Builds training data pipelines.
- data.py → ImageNet data: loads · Instantiates dataset sources.
- data.py → Corruption base: applies · Routes samples through a degradation interface.
- Corruption base → Corruption ops: implements · Concrete corruption operators.
- train.py → UNet: trains · Optimizes the pixel-space model.
- train.py → Diffusion core: uses · Relies on diffusion process definitions.
- train.py → Guided losses: optimizes · Uses diffusion losses and objectives.
- UNet → Guided utils: relies on · Uses attention, fp16, and helper utilities.
- Diffusion core → Guided utils: relies on · Shares scheduling and runtime helpers.
- sample.py → Checkpoints: loads · Consumes pretrained model artifacts.
- sample.py → UNet: runs · Generates restorations with the pixel model.
- sample.py → Runner: runs · Can invoke the alternative i2sb pipeline.
- Runner → i2sb diffusion: orchestrates · Coordinates the i2sb diffusion process.
- Runner → Network: orchestrates · Builds and executes the model wrapper.
- Network → VQGAN: encodes · Uses VQGAN for latent compression and decoding.
- VQGAN → Taming: depends on · Uses vendored VQGAN/taming components.
- compute_metrices.py → Evaluation pkg: uses · Computes FID and related scores.
- paper_vis.py → sample.py: consumes · Visualizations are built from sampled outputs.
- train.py → Logging: reports to · Writes training progress and metrics.
- train.py → Distributed: scales with · Supports distributed execution.
- Runner → Logging: reports to · Logs i2sb training/inference status.
- Runner → Distributed: scales with · Supports multi-device execution.
- Evaluation pkg → paper_vis.py: feeds · Evaluation outputs support figures and comparisons.