# 6ungying/piff-master02 architecture

> Architecture diagram of the GitHub repository 6ungying/piff-master02, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-05-17). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/6ungying/piff-master02
- Repository: https://github.com/6ungying/piff-master02

## Overview

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`](https://github.com/6ungying/piff-master02/blob/master/train.py) (train.py), [`i2sb/runner.py`](https://github.com/6ungying/piff-master02/blob/master/i2sb/runner.py) (Runner), [`sample.py`](https://github.com/6ungying/piff-master02/blob/master/sample.py) (sample.py), [`data.py`](https://github.com/6ungying/piff-master02/blob/master/data.py) (data.py)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entrypoints["Entrypoints"]
  node_train_py["train.py<br/>training entrypoint<br/>[train.py]"]
  node_sample_py["sample.py<br/>inference entrypoint<br/>[sample.py]"]
  node_metrics_py["compute_metrices.py<br/>evaluation entrypoint"]
  node_paper_vis_py["paper_vis.py<br/>visualization entrypoint<br/>[paper_vis.py]"]
  node_configs["configs<br/>experiment config"]
end

subgraph group_data["Data & Degradation"]
  node_data_py["data.py<br/>data loader<br/>[data.py]"]
  node_dataset_imagenet["ImageNet data<br/>dataset module<br/>[imagenet.py]"]
  node_corruption_base["Corruption base<br/>degradation interface<br/>[base.py]"]
  node_corruption_ops["Corruption ops<br/>degradation models"]
end

subgraph group_pixel["Pixel Diffusion"]
  node_guided_unet["UNet<br/>diffusion backbone<br/>[unet.py]"]
  node_guided_diffusion_core["Diffusion core<br/>diffusion process"]
  node_guided_losses["Guided losses<br/>objective module<br/>[losses.py]"]
  node_guided_util["Guided utils<br/>model utilities"]
end

subgraph group_latent["Latent Diffusion"]
  node_i2sb_runner["Runner<br/>training orchestration<br/>[runner.py]"]
  node_i2sb_diffusion["i2sb diffusion<br/>diffusion stack<br/>[diffusion.py]"]
  node_i2sb_network["Network<br/>model wrapper<br/>[network.py]"]
  node_vqgan["VQGAN<br/>latent codec<br/>[vqgan.py]"]
  node_taming["Taming<br/>vendored components"]
end

subgraph group_eval["Evaluation"]
  node_evaluation["Evaluation pkg<br/>metric stack"]
end

subgraph group_infra["Infrastructure"]
  node_logger["Logging<br/>runtime support<br/>[logger.py]"]
  node_distributed_util["Distributed<br/>runtime support"]
  node_checkpoint_json[("Checkpoints<br/>artifact manifest<br/>[checkpoint.json]")]
end

node_train_py -->|"reads"| node_configs
node_train_py -->|"uses"| node_data_py
node_data_py -->|"loads"| node_dataset_imagenet
node_data_py -->|"applies"| node_corruption_base
node_corruption_base -->|"implements"| node_corruption_ops
node_train_py -->|"trains"| node_guided_unet
node_train_py -->|"uses"| node_guided_diffusion_core
node_train_py -->|"optimizes"| node_guided_losses
node_guided_unet -->|"relies on"| node_guided_util
node_guided_diffusion_core -->|"relies on"| node_guided_util
node_sample_py -->|"loads"| node_checkpoint_json
node_sample_py -->|"runs"| node_guided_unet
node_sample_py -->|"runs"| node_i2sb_runner
node_i2sb_runner -->|"orchestrates"| node_i2sb_diffusion
node_i2sb_runner -->|"orchestrates"| node_i2sb_network
node_i2sb_network -->|"encodes"| node_vqgan
node_vqgan -->|"depends on"| node_taming
node_metrics_py -->|"uses"| node_evaluation
node_paper_vis_py -->|"consumes"| node_sample_py
node_train_py -->|"reports to"| node_logger
node_train_py -->|"scales with"| node_distributed_util
node_i2sb_runner -->|"reports to"| node_logger
node_i2sb_runner -->|"scales with"| node_distributed_util
node_evaluation -->|"feeds"| node_paper_vis_py

click node_train_py "https://github.com/6ungying/piff-master02/blob/master/train.py"
click node_sample_py "https://github.com/6ungying/piff-master02/blob/master/sample.py"
click node_metrics_py "https://github.com/6ungying/piff-master02/blob/master/compute_metrices.py"
click node_paper_vis_py "https://github.com/6ungying/piff-master02/blob/master/paper_vis.py"
click node_configs "https://github.com/6ungying/piff-master02/tree/master/configs"
click node_data_py "https://github.com/6ungying/piff-master02/blob/master/data.py"
click node_dataset_imagenet "https://github.com/6ungying/piff-master02/blob/master/dataset/imagenet.py"
click node_corruption_base "https://github.com/6ungying/piff-master02/blob/master/corruption/base.py"
click node_corruption_ops "https://github.com/6ungying/piff-master02/tree/master/corruption"
click node_guided_unet "https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/unet.py"
click node_guided_diffusion_core "https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/gaussian_diffusion.py"
click node_guided_losses "https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/losses.py"
click node_guided_util "https://github.com/6ungying/piff-master02/tree/master/guided_diffusion"
click node_i2sb_runner "https://github.com/6ungying/piff-master02/blob/master/i2sb/runner.py"
click node_i2sb_diffusion "https://github.com/6ungying/piff-master02/blob/master/i2sb/diffusion.py"
click node_i2sb_network "https://github.com/6ungying/piff-master02/blob/master/i2sb/network.py"
click node_vqgan "https://github.com/6ungying/piff-master02/blob/master/i2sb/VQGAN/vqgan.py"
click node_taming "https://github.com/6ungying/piff-master02/tree/master/i2sb/VQGAN/taming"
click node_evaluation "https://github.com/6ungying/piff-master02/tree/master/evaluation"
click node_logger "https://github.com/6ungying/piff-master02/blob/master/logger.py"
click node_distributed_util "https://github.com/6ungying/piff-master02/blob/master/distributed_util.py"
click node_checkpoint_json "https://github.com/6ungying/piff-master02/blob/master/checkpoint.json"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train_py,node_sample_py,node_metrics_py,node_paper_vis_py,node_configs toneBlue
class node_data_py,node_dataset_imagenet,node_corruption_base,node_corruption_ops toneAmber
class node_guided_unet,node_guided_diffusion_core,node_guided_losses,node_guided_util toneMint
class node_i2sb_runner,node_i2sb_diffusion,node_i2sb_network,node_vqgan,node_taming toneRose
class node_evaluation toneIndigo
class node_logger,node_distributed_util,node_checkpoint_json toneTeal
```

## Components

### Entrypoints

Training, sampling, evaluation, and visualization scripts.

- [train.py](https://github.com/6ungying/piff-master02/blob/master/train.py): training entrypoint · Main model fitting script. · `train.py`
- [sample.py](https://github.com/6ungying/piff-master02/blob/master/sample.py): inference entrypoint · Generation and restoration sampling script. · `sample.py`
- [compute_metrices.py](https://github.com/6ungying/piff-master02/blob/master/compute_metrices.py): evaluation entrypoint · Runs quantitative evaluation over outputs. · `compute_metrices.py`
- [paper_vis.py](https://github.com/6ungying/piff-master02/blob/master/paper_vis.py): visualization entrypoint · Produces figures and visualization artifacts. · `paper_vis.py`
- [configs](https://github.com/6ungying/piff-master02/tree/master/configs): experiment config · YAML templates for model variants, latent sizes, and degradation setups. · `configs`

### Data & Degradation

Dataset loading and image corruption pipelines.

- [data.py](https://github.com/6ungying/piff-master02/blob/master/data.py): data loader · Top-level dataset and dataloader orchestration. · `data.py`
- [ImageNet data](https://github.com/6ungying/piff-master02/blob/master/dataset/imagenet.py): dataset module · ImageNet-centric dataset implementation and validation split metadata. · `dataset/imagenet.py`
- [Corruption base](https://github.com/6ungying/piff-master02/blob/master/corruption/base.py): degradation interface · Shared interface for forward image corruption models. · `corruption/base.py`
- [Corruption ops](https://github.com/6ungying/piff-master02/tree/master/corruption): degradation models · Blur, inpainting, JPEG, super-resolution, and mixtures. · `corruption`

### Pixel Diffusion

Core guided diffusion stack for image restoration.

- [UNet](https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/unet.py): diffusion backbone · Main neural network architecture for pixel-space diffusion. · `guided_diffusion/unet.py`
- [Diffusion core](https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/gaussian_diffusion.py): diffusion process · Gaussian diffusion, respacing, and sampling math. · `guided_diffusion/gaussian_diffusion.py`
- [Guided losses](https://github.com/6ungying/piff-master02/blob/master/guided_diffusion/losses.py): objective module · Loss functions and training objectives for diffusion models. · `guided_diffusion/losses.py`
- [Guided utils](https://github.com/6ungying/piff-master02/tree/master/guided_diffusion): model utilities · Attention, fp16, logging, script helpers, and runtime utilities. · `guided_diffusion`

### Latent Diffusion

I2SB and VQGAN-based latent model stack.

- [Runner](https://github.com/6ungying/piff-master02/blob/master/i2sb/runner.py): training orchestration · Higher-level training and inference loop for the i2sb stack. · `i2sb/runner.py`
- [i2sb diffusion](https://github.com/6ungying/piff-master02/blob/master/i2sb/diffusion.py): diffusion stack · Alternative diffusion modules and process definitions. · `i2sb/diffusion.py`
- [Network](https://github.com/6ungying/piff-master02/blob/master/i2sb/network.py): model wrapper · Model wiring, embeddings, and network assembly. · `i2sb/network.py`
- [VQGAN](https://github.com/6ungying/piff-master02/blob/master/i2sb/VQGAN/vqgan.py): latent codec · Latent compression and reconstruction model used for bottleneck diffusion. · `i2sb/VQGAN/vqgan.py`
- [Taming](https://github.com/6ungying/piff-master02/tree/master/i2sb/VQGAN/taming): vendored components · Vendored VQGAN datasets, modules, and transformer utilities. · `i2sb/VQGAN/taming`

### Evaluation

Metrics and scoring for generated outputs.

- [Evaluation pkg](https://github.com/6ungying/piff-master02/tree/master/evaluation): metric stack · FID utilities and ResNet-based evaluator. · `evaluation`

### Infrastructure

Logging, distributed execution, and utility support.

- [Logging](https://github.com/6ungying/piff-master02/blob/master/logger.py): runtime support · Experiment logging and status reporting. · `logger.py`
- [Distributed](https://github.com/6ungying/piff-master02/blob/master/distributed_util.py): runtime support · Multi-device execution helpers. · `distributed_util.py`
- [Checkpoints](https://github.com/6ungying/piff-master02/blob/master/checkpoint.json): artifact manifest · External pretrained weights and checkpoint references. · `checkpoint.json`

## Connections

- 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.