---
license: apache-2.0
tags:
  - diffusion
  - text-to-image
  - safety
  - dose-response
base_model: Photoroom/PRX
datasets:
  - lehduong/flux_generated
  - LucasFang/FLUX-Reason-6M
  - brivangl/midjourney-v6-llava
pipeline_tag: text-to-image
---

# Dose-Response C4: Original proportion (~1.21% unsafe), 1M scale

This model is part of a **dose-response experiment** studying how the fraction of unsafe content in training data affects the safety of generated images from text-to-image diffusion models.

## Model Details

| | |
|---|---|
| **Architecture** | PRX-1.2B (Photoroom diffusion model) |
| **Parameters** | 1.2B (denoiser only) |
| **Resolution** | 512px |
| **Condition** | C4 — Original proportion (~1.21% unsafe), 1M scale |
| **Unsafe fraction** | ~1.21% (original) |
| **Training set size** | 1M images |
| **Training steps** | 100K batches |
| **Batch size** | 1024 (global) |
| **Precision** | bf16 |
| **Hardware** | 8x H200 GPUs |

## Condition Description

Downscaled to 1M images while preserving the original ~1.21% unsafe proportion (12K unsafe, 988K safe).

## Dose-Response Conditions Overview

This model is one of 7 conditions in the dose-response experiment:

| Condition | Unsafe Fraction | Dataset Scale | Description |
|-----------|----------------|---------------|-------------|
| **C0** | 0% | Full (~7.85M) | All unsafe removed |
| **C1** | 5% | Full (~8.24M) | Unsafe oversampled to 5% |
| **C2** | 10% | Full (~8.72M) | Unsafe oversampled to 10% |
| **C3** | ~1.21% | Full (~7.94M) | Original composition |
| **C4** | ~1.21% | 1M | Original proportion, downscaled |
| **C5** | ~9.6% | 1M | All unsafe included, downscaled |
| **C6** | ~1.21% | 100K | Original proportion, small scale |

## Training Details

- **Base architecture**: [PRX](https://github.com/Photoroom/PRX) 1.2B
- **Text encoder**: T5-Gemma-2B (frozen)
- **VAE**: Identity (no compression)
- **Optimizer**: Muon
- **Algorithms**: TREAD + REPA-v3 + LPIPS + Perceptual DINO + EMA
- **EMA smoothing**: 0.999 (updated every 10 batches)
- **Training data sources**: `lehduong/flux_generated`, `LucasFang/FLUX-Reason-6M`, `brivangl/midjourney-v6-llava`
- **Safety annotations**: Training data annotated with [LlavaGuard-7B](https://huggingface.co/AIML-TUDA/LlavaGuard-v1.2-7B-OV) to classify images as safe/unsafe

## Files

- `denoiser.pt` — Consolidated single-file checkpoint (EMA weights, ready for inference)
- `distributed/` — Original FSDP distributed checkpoint shards
- `config.yaml` — Full Hydra training configuration

## Usage

```python
import torch

# Load consolidated checkpoint
state_dict = torch.load("denoiser.pt", map_location="cpu")
# Keys are in format: denoiser.*
```

For the full generation pipeline, see the [diffusion_safety](https://github.com/felifri/diffusion_safety) repository.

## Citation

If you use these models, please cite the associated paper and the PRX architecture.

## License

Apache 2.0