# wahib-elbessali/conditional-stylegan2-ada architecture

> Architecture diagram of the GitHub repository wahib-elbessali/conditional-stylegan2-ada, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/wahib-elbessali/conditional-stylegan2-ada
- Repository: https://github.com/wahib-elbessali/conditional-stylegan2-ada

## Overview

This repository trains class-conditional StyleGAN2-ADA models for brain MRI, chest X-ray, and dementia MRI, serves labeled synthetic images through a Gradio demo, and evaluates image quality and classifier utility. The main workflow is dataset preparation, per-domain training or fine-tuning, periodic cross-worker weight averaging, checkpoint use for generation, and evaluation against real images. Training, evaluation, and the demo are separate entry points. Training and evaluation excerpts are partial, so some links reflect the README’s explicit workflow rather than inspected implementation details.

Main entry points: [`distributed_training.py`](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py) (Training pipeline), [`evaluation.py`](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py) (Evaluation pipeline), [`app.py`](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py) (Gradio interface)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["Distributed training"]
  node_trainer["Training pipeline"]
  node_dataset["Dataset loaders"]
  node_sampler["Stratified sampler"]
  node_gan["Conditional GAN"]
  node_ada["ADA augmentation"]
  node_fine_tune["Domain fine-tuning"]
  node_sync["Weight averaging"]
  node_checkpoints["Checkpoints"]
end

subgraph group_demo["Image generation"]
  node_ui["Gradio interface<br/>[app.py]"]
  node_loader["Checkpoint loader<br/>[app.py]"]
  node_generator["Image generator<br/>[app.py]"]
  node_gallery["Generated images<br/>[app.py]"]
end

subgraph group_evaluation["Model evaluation"]
  node_eval_runner["Evaluation pipeline<br/>[evaluation.py]"]
  node_metric_eval["Quality metrics<br/>[evaluation.py]"]
  node_classifier_eval["Classifier utility<br/>[evaluation.py]"]
  node_results["Evaluation results<br/>[evaluation.py]"]
end

node_user(("Demo user"))
node_kaggle[("Kaggle datasets")]
node_hub[("HuggingFace Hub")]
node_workers(("GPU workers"))

node_kaggle -->|"provides data"| node_dataset
node_trainer -->|"loads samples"| node_dataset
node_dataset -->|"labels batches"| node_sampler
node_trainer -->|"uses batches"| node_sampler
node_trainer -->|"trains models"| node_gan
node_trainer -->|"updates augmentation"| node_ada
node_trainer -->|"runs phases"| node_fine_tune
node_fine_tune -->|"loads base weights"| node_checkpoints
node_trainer -->|"coordinates sync"| node_sync
node_workers -->|"participate in"| node_sync
node_sync -->|"exchanges weights"| node_hub
node_trainer -->|"saves weights"| node_checkpoints
node_trainer -->|"stores checkpoints"| node_hub
node_user -->|"selects condition"| node_ui
node_ui -->|"requests model"| node_loader
node_loader -->|"downloads checkpoint"| node_hub
node_loader -->|"loads weights"| node_generator
node_ui -->|"requests images"| node_generator
node_generator -->|"returns images"| node_gallery
node_eval_runner -->|"loads real data"| node_dataset
node_eval_runner -->|"downloads checkpoints"| node_hub
node_eval_runner -->|"generates samples"| node_gan
node_eval_runner -->|"computes FID KID IS"| node_metric_eval
node_eval_runner -->|"tests utility"| node_classifier_eval
node_metric_eval -->|"records scores"| node_results
node_classifier_eval -->|"records scores"| node_results

click node_trainer "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_dataset "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_sampler "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_gan "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_ada "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_fine_tune "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_sync "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_checkpoints "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py"
click node_ui "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py"
click node_loader "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py"
click node_generator "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py"
click node_gallery "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py"
click node_eval_runner "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py"
click node_metric_eval "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py"
click node_classifier_eval "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py"
click node_results "https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py"

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_trainer,node_dataset,node_sampler,node_gan,node_ada,node_fine_tune,node_sync,node_checkpoints,node_user toneBlue
class node_ui,node_loader,node_generator,node_gallery,node_kaggle,node_hub toneAmber
class node_eval_runner,node_metric_eval,node_classifier_eval,node_results toneMint
class node_workers toneRose
```

## Components

### Distributed training

- [Training pipeline](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Dataset loaders](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Stratified sampler](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Conditional GAN](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [ADA augmentation](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Domain fine-tuning](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Weight averaging](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`
- [Checkpoints](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/distributed_training.py): `distributed_training.py`

### Image generation

- [Gradio interface](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py): `app.py`
- [Checkpoint loader](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py): `app.py`
- [Image generator](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py): `app.py`
- [Generated images](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/app.py): `app.py`

### Model evaluation

- [Evaluation pipeline](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py): `evaluation.py`
- [Quality metrics](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py): `evaluation.py`
- [Classifier utility](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py): `evaluation.py`
- [Evaluation results](https://github.com/wahib-elbessali/conditional-stylegan2-ada/blob/main/evaluation.py): `evaluation.py`

### Other components

- **Demo user**
- **Kaggle datasets**
- **HuggingFace Hub**
- **GPU workers**

## Connections

- Kaggle datasets → Dataset loaders: provides data
- Training pipeline → Dataset loaders: loads samples
- Dataset loaders → Stratified sampler: labels batches
- Training pipeline → Stratified sampler: uses batches
- Training pipeline → Conditional GAN: trains models
- Training pipeline → ADA augmentation: updates augmentation
- Training pipeline → Domain fine-tuning: runs phases
- Domain fine-tuning → Checkpoints: loads base weights
- Training pipeline → Weight averaging: coordinates sync
- GPU workers → Weight averaging: participate in
- Weight averaging → HuggingFace Hub: exchanges weights
- Training pipeline → Checkpoints: saves weights
- Training pipeline → HuggingFace Hub: stores checkpoints
- Demo user → Gradio interface: selects condition
- Gradio interface → Checkpoint loader: requests model
- Checkpoint loader → HuggingFace Hub: downloads checkpoint
- Checkpoint loader → Image generator: loads weights
- Gradio interface → Image generator: requests images
- Image generator → Generated images: returns images
- Evaluation pipeline → Dataset loaders: loads real data
- Evaluation pipeline → HuggingFace Hub: downloads checkpoints
- Evaluation pipeline → Conditional GAN: generates samples
- Evaluation pipeline → Quality metrics: computes FID KID IS
- Evaluation pipeline → Classifier utility: tests utility
- Quality metrics → Evaluation results: records scores
- Classifier utility → Evaluation results: records scores