fluxflow-training 0.3.1

Training tools and scripts for FluxFlow text-to-image generation

FluxFlow Training

Training tools and scripts for FluxFlow text-to-image generation models.

Installation

Production Install

pip install fluxflow-training

What gets installed:

• fluxflow-training - Training scripts and configuration tools
• fluxflow core package (automatically installed as dependency)
• CLI commands: fluxflow-train, fluxflow-generate

Development Install

git clone https://github.com/danny-mio/fluxflow-training.git
cd fluxflow-training
pip install -e ".[dev]"

---

🚧 Training Status

Models Currently In Training: FluxFlow is actively training models following the systematic TRAINING_VALIDATION_PLAN.md.

Current Phase: Phase 1 - VAE Training (Weeks 1-4)

Progress:

• 🔄 Bezier VAE training in progress
• ⏳ ReLU baseline VAE pending
• ⏳ Flow models pending VAE completion

When Available: Trained checkpoints and empirical performance metrics will be published to MODEL_ZOO.md upon validation completion.

Note: All performance claims in documentation are theoretical targets pending empirical validation.

---

Hardware Requirements for Training

Current Setup (December 2025 validation):

• GPU: NVIDIA A6000 (48GB VRAM) via Paperspace
• Constraint: Service interrupts every 6 hours
• Solution: Automatic checkpoint/resume (every 30-60 min)

Alternative Options:

• Local: 1× RTX 4090 (24GB) or 2× RTX 3090 (24GB each)
• Cloud: AWS p3.2xlarge (V100 16GB), GCP A100 (40GB)

Memory Requirements by Training Mode (empirical measurements, Dec 2025):

• VAE Training (batch_size=4, vae_dim=128, img_size=1024):
  • Without GAN: ~18-22GB VRAM
  • With GAN + LPIPS: ~28-35GB VRAM
  • With GAN + LPIPS + SPADE: ~35-42GB VRAM
  • Peak observed: 47.4GB on A6000 48GB (triggered OOM, requires optimization)
• Flow Training (batch_size=1, feature_maps_dim=128):
  • ~24-30GB VRAM
• Minimum viable (reduced dimensions, smaller images):
  • 16GB VRAM for VAE (batch_size=2, vae_dim=64, img_size=512)
  • 24GB VRAM for Flow (batch_size=1, feature_maps_dim=64)

OOM Prevention (if you hit 47GB+ on 48GB GPU):

• Reduce batch size: batch_size: 2 or 1
• Disable LPIPS: use_lpips: false (saves ~6-8GB)
• Reduce image size: img_size: 512 (saves ~10-15GB)
• Use FP16 (if supported): use_fp16: true (saves ~20-30%)
• See TRAINING_GUIDE.md "Limited VRAM Strategy"

Cost Comparison:

Platform	GPU	$/hr	Est. Total (500 hrs)
Paperspace	A6000 48GB	$0.76	~$456 ✓
AWS	p3.2xlarge (V100 16GB)	$3.06	~$1,530
GCP	A100 40GB	$3.67	~$1,835
Lambda Labs	A100 40GB	$1.10	~$550

---

Pre-download LPIPS Weights (Optional)

The training uses LPIPS for perceptual loss, which requires VGG16 weights (~528MB). To pre-download:

python -c "import lpips; lpips.LPIPS(net='vgg')"

Weights will be cached in ~/.cache/torch/hub/checkpoints/. If not pre-downloaded, they'll download automatically on first training run.

Features

Core Training Capabilities

• 🎯 Pipeline Training Mode (v0.2.0+, FULLY IMPLEMENTED)

  • Multi-step sequential training with independent configs per step
  • Per-step freeze/unfreeze of model components
  • Loss-threshold transitions with early stopping
  • Full checkpoint resume from any step/epoch/batch
  • 1035 lines in pipeline_orchestrator.py
  • See PIPELINE_ARCHITECTURE.md

• 🎨 GAN-Only Training Mode (v0.2.0+, FULLY IMPLEMENTED)

  • Train encoder/decoder with adversarial loss only (no reconstruction)
  • Spatial conditioning (SPADE) without pixel-perfect reconstruction
  • Faster training with focused gradient flow
  • Example config: see PIPELINE_ARCHITECTURE.md "GAN-Only Mode" section

• VAE Training

  • Variational autoencoders with GAN losses
  • LPIPS perceptual loss support (adds ~6-8GB VRAM)
  • SPADE spatial conditioning
  • KL divergence with beta warmup and free bits

• Flow Training

  • Flow-based diffusion transformers
  • Text-to-image generation with classifier-free guidance (CFG)
  • Industry-standard training approach (same as Stable Diffusion, Flux.1)

• ✨ Classifier-Free Guidance (CFG) (v0.3.0+)

  • Train models to generate with or without text conditioning
  • Single model learns both conditional p(x|text) and unconditional p(x)
  • Inference-time guidance scale control (1.0-15.0, recommended: 3.0-7.0)
  • Enables quality/diversity trade-off without retraining
  • Proven approach: Stable Diffusion, DALL-E 2, Imagen, Flux.1
  • See config.cfg.example.yaml for setup

Training Infrastructure

• Data Loading: Efficient dataset handling with WebDataset support for large-scale training
• Checkpointing: Robust checkpoint management with automatic resume from interruptions
• Training Visualization: Automatic diagram generation for loss curves and metrics
• Optimizers: Multiple optimizer support (AdamW, Adam, SGD, RMSprop)
• Schedulers: Various learning rate schedulers (CosineAnnealing, StepLR, ExponentialLR, ReduceLROnPlateau)
• Mixed Precision: Accelerate training with automatic mixed precision (fp16)

Monitoring & Logging

• Enhanced Console Output (v0.2.0)

  • Real-time batch timing (seconds/batch)
  • Comprehensive loss display (VAE, KL, G, D, LPIPS)
  • Step/epoch/batch progress tracking

• Step-Specific Metrics (Pipeline Mode)

  • Separate JSONL files per training step
  • Automatic graph generation per step
  • Full training history preservation

• Mid-Epoch Sample Generation (v0.2.0)

  • Samples generated at configurable intervals
  • Structured naming with step/epoch/batch numbers
  • Prevents overwrites with unique identifiers

Quick Start

Basic Training

# Create a config file (see config.example.yaml)
fluxflow-train --config config.yaml

# With automatic diagram generation
fluxflow-train --config config.yaml --generate_diagrams

Pipeline Training (NEW in v0.2.0)

Multi-step training for hypothesis testing and staged optimization:

training:
  pipeline:
    steps:
      # Step 1: VAE warmup
      - name: "vae_warmup"
        n_epochs: 10
        train_vae: true
        gan_training: false
        freeze:
          - flow_processor
          - text_encoder
      
      # Step 2: Add GAN
      - name: "vae_gan"
        n_epochs: 20
        train_vae: true
        gan_training: true
        use_lpips: true
        
      # Step 3: Flow training
      - name: "flow"
        n_epochs: 30
        train_vae: false      # Freeze VAE
        train_diff: true      # Train flow

Run:

fluxflow-train --config pipeline_config.yaml

Validate config before training:

fluxflow-train --config pipeline_config.yaml --validate-pipeline

See PIPELINE_ARCHITECTURE.md for complete documentation.

GAN-Only Training (NEW in v0.2.0)

Train with adversarial loss only (no pixel-level reconstruction):

training:
  pipeline:
    steps:
      - name: "gan_only"
        n_epochs: 20
        train_vae: false          # Skip reconstruction loss
        gan_training: true        # Train GAN discriminator
        train_spade: true         # With spatial conditioning
        use_lpips: false

Use cases:

• Spatial structure learning without pixel constraints
• Faster training (no reconstruction computation)
• SPADE conditioning experiments
• Saves ~8-12GB VRAM by skipping reconstruction loss computation

Classifier-Free Guidance Training (NEW in v0.3.0)

Train text-to-image models with CFG for superior prompt adherence:

training:
  pipeline:
    steps:
      # Stage 1-2: VAE training (see config.cfg.example.yaml)
      
      # Stage 3: Flow with CFG ✨
      - name: "flow_cfg"
        n_epochs: 100
        train_diff: true
        cfg_dropout_prob: 0.10  # 10% null conditioning (common setting)
        use_ema: true
        
        freeze:
          - compressor  # Freeze VAE
          - expander

Generate with CFG:

fluxflow-generate \
    --model_checkpoint outputs/flux_cfg/final.safetensors \
    --text_prompts_path prompts/ \
    --output_path outputs/ \
    --use_cfg \
    --guidance_scale 5.0  # Range: 1.0-15.0, recommended: 3.0-7.0

Guidance scale examples (tune for your model):

• 1.0: Standard conditional (no guidance)
• 3.0-7.0: Moderate guidance (RECOMMENDED - balanced quality/creativity)
• 7.0-15.0: Strong guidance (may oversaturate or lose diversity)

See config.cfg.example.yaml for complete CFG training example.

Generating Images

# Create a directory with .txt files containing prompts
mkdir prompts
echo "a beautiful sunset over mountains" > prompts/sunset.txt

# Generate images (standard)
fluxflow-generate \
    --model_checkpoint path/to/checkpoint.safetensors \
    --text_prompts_path prompts/ \
    --output_path outputs/

# Generate with CFG (if model trained with cfg_dropout_prob > 0)
fluxflow-generate \
    --model_checkpoint path/to/checkpoint.safetensors \
    --text_prompts_path prompts/ \
    --output_path outputs/ \
    --use_cfg \
    --guidance_scale 5.0

Visualizing Training Progress

# Training metrics are automatically logged to outputs/graph/training_metrics.jsonl
# In pipeline mode: outputs/graph/training_metrics_{stepname}.jsonl

# Generate diagrams from logged metrics
python src/fluxflow_training/scripts/generate_training_graphs.py outputs/

# Diagrams are saved to outputs/graph/:
# - training_losses.png (VAE, Flow, Discriminator, Generator, LPIPS)
# - kl_loss.png (KL divergence with beta warmup)
# - learning_rates.png (LR schedules)
# - batch_times.png (training speed)
# - training_overview.png (combined overview)
# - training_summary.txt (statistics)

Console Output Examples

Pipeline Training Mode

PIPELINE STEP 1/3: vae_warmup
Description: Warmup VAE without GAN
Duration: 10 epochs
================================================================================

[00:15:23] Step vae_warmup (1/3) | Epoch 5/10 | Batch 127/500 | VAE: 0.0234 | KL: 12.45 | 3.2s/batch

GAN Training Mode

[00:15:23] Step vae_gan (2/3) | Epoch 5/20 | Batch 127/500 | VAE: 0.0234 | KL: 12.45 | G: 0.156 | D: 0.089 | LPIPS: 0.0812 | 3.2s/batch

Metrics Shown:

• VAE - Reconstruction loss (L1 + MSE)
• KL - KL divergence
• G - GAN generator loss (when gan_training: true)
• D - GAN discriminator loss (when gan_training: true)
• LPIPS - Perceptual loss (when use_lpips: true)
• Xs/batch - Average batch processing time

Sample Naming Convention

Pipeline Mode (v0.2.0)

Mid-epoch samples:

vae_warmup_001_005_00127_abc123-original.webp
{stepname}_{step}_{epoch}_{batch}_{hash}-{suffix}.webp

End-of-epoch samples:

vae_warmup_001_005_abc123-original.webp
{stepname}_{step}_{epoch}_{hash}-{suffix}.webp

Suffixes:

• -original or _ns_i - Input image
• -nr_o - VAE reconstruction (no random sampling)
• _ns_o - VAE reconstruction (with sampling)
• -ctx - Context vector visualization
• -nc - Non-context visualization

Package Contents

• fluxflow_training.training - Training logic and trainers

  • pipeline_orchestrator.py - Multi-step pipeline execution (1035 lines)
  • pipeline_config.py - Pipeline configuration and validation
  • vae_trainer.py - VAE/GAN training logic
  • flow_trainer.py - Flow model training
  • checkpoint_manager.py - Checkpoint save/resume

• fluxflow_training.data - Dataset implementations and transforms

  • datasets.py - Image/caption datasets with WebDataset support
  • transforms.py - Data augmentation pipelines

• fluxflow_training.scripts - CLI scripts for training and generation

  • train.py - Main training script (legacy + pipeline mode)
  • generate_training_graphs.py - Visualization generation

Configuration

Training is configured via YAML files. See docs/TRAINING_GUIDE.md for detailed configuration options.

Basic Configuration Example

model:
  vae_dim: 128
  feature_maps_dim: 128
  text_embedding_dim: 1024

data:
  data_path: "/path/to/images"
  captions_file: "/path/to/captions.txt"  # image_name<tab>caption
  fixed_prompt_prefix: null  # Optional: e.g., "style anime" to prepend to all prompts
  img_size: 1024
  reduced_min_sizes: null  # Optional: [128, 256, 512]

training:
  n_epochs: 100
  batch_size: 1
  lr: 0.00001
  train_vae: true
  gan_training: true
  use_lpips: true
  use_fp16: false

output:
  output_path: "outputs/flux"
  log_interval: 10
  checkpoint_save_interval: 50
  samples_per_checkpoint: 1  # Generate samples at each checkpoint
  sample_sizes:  # Optional: generate samples at various sizes
    - 512
    - [768, 512]  # landscape
    - 1024

Pipeline Configuration Example

See test_pipeline_minimal.yaml or config.example.yaml for complete examples.

Key difference: Pipeline configs use training.pipeline.steps instead of flat training config.

training:
  batch_size: 4
  workers: 8
  
  pipeline:
    steps:
      - name: "step1"
        n_epochs: 10
        train_vae: true
        # ... step-specific config

Documentation

Comprehensive Guides

• PIPELINE_ARCHITECTURE.md - Pipeline training mode documentation

  • Configuration reference
  • Checkpoint format
  • Sample naming conventions
  • Troubleshooting
  • Status: ✅ FULLY IMPLEMENTED (1035 lines in pipeline_orchestrator.py)

• TRAINING_GUIDE.md - Complete training guide

  • Detailed configuration options
  • Dataset preparation
  • Training strategies
  • Memory optimization strategies
  • Best practices

• CONTRIBUTING.md - Development guide

  • Setting up development environment
  • Running tests
  • Code style guidelines
  • Pipeline config testing

Example Configurations

• config.example.yaml - Complete configuration template with all options
• test_pipeline_minimal.yaml - Minimal working pipeline example
• config.example.sh - Environment setup script

What's New in v0.2.1

Critical Optimizations (Dec 2025)

• Memory optimizations to prevent OOM on 48GB GPUs
  • Removed LPIPS gradient checkpointing (caused OOM at 47.4GB)
  • Removed dataloader prefetch_factor (caused memory spikes)
  • CUDA cache clearing between batches
  • See CHANGELOG.md for full details

Major Features (v0.2.0)

• Pipeline Training Mode - Multi-step sequential training with per-step configs
• GAN-Only Mode - Train with adversarial loss only (no reconstruction)
• Enhanced Logging - Batch timing, step-specific metrics, correct GAN loss display
• Mid-Epoch Samples - Configurable sample generation with batch numbers

Bug Fixes

• Fixed encoder gradient flow in GAN-only mode
• Fixed GAN loss logging (generator/discriminator keys)
• Fixed EMA creation for GAN-only mode
• Fixed metrics logging to respect training modes
• Fixed sample generation to prevent overwrites
• Fixed R1 penalty gradient computation

Breaking Changes

Config format: Pipeline configs must use steps: wrapper:

# OLD (will error)
training:
  pipeline:
    - name: "step1"

# NEW (required)
training:
  pipeline:
    steps:
      - name: "step1"

See CHANGELOG.md for complete version history.

Links

• GitHub Repository
• Core Package
• Web UI
• ComfyUI Plugin
• PyPI Package (coming soon - not yet published)

License

MIT License - see LICENSE file for details.