Upload folder using huggingface_hub

README.md

@@ -41,18 +41,31 @@ out = model.generate(
     "Rain falls heavily on a metal roof with distant thunder",
     num_inference_steps=50,
     guidance_scale=4.0,
-    generator=torch.Generator(device=model.device).manual_seed(0),
+    seed=0,  # same default as the original audio_dit/sample.py
 )
 
 sf.write("rain.wav", out.audios[0], out.sampling_rate)
 ```
 
+To reproduce an existing `audio_dit/sample.py` result, use the **same prompt,
+steps, CFG scale and seed**. For example, the original high-guidance test was:
+
+```python
+out = model.generate(
+    "Rain falls heavily on a roof top with distant thunder",
+    num_inference_steps=200,
+    guidance_scale=16.0,
+    seed=0,
+)
+```
+
 ### Useful `generate` knobs
 
 | Argument | Default | Meaning |
 |---|---|---|
 | `num_inference_steps` | 50 | DDIM steps |
 | `guidance_scale` | 4.0 | CFG scale (try 5–7 for crisper AudioCaps-style prompts) |
+| `seed` | 0 | Initial latent noise seed; set it for reproducible A/B tests |
 | `output_type` | `"np"` | `"np"` / `"pt"` / `"latent"` |
 | `generator` | `None` | `torch.Generator` for reproducibility |
 
@@ -63,15 +76,32 @@ Prompts work best in AudioCaps style: short concrete sound events.
 From the research repo:
 
 ```bash
+# DDP checkpoints already embed train meta — no cache needed
 python qadit_hf/convert_checkpoint.py \
   --ckpt path/to/ckpt_0023999.pt \
-  --cache path/to/precompute_cache \
   --out ./qadit-export \
   --push-to-hub USER/qadit
+
+# Or pass the known scale explicitly (from your training log, e.g. 0.9503)
+python qadit_hf/convert_checkpoint.py \
+  --ckpt path/to/ckpt_0023999.pt \
+  --latent-scale 0.9503 \
+  --out ./qadit-export
 ```
 
 This writes `config.json`, `model.safetensors`, and the remote-code `.py` files.
-`latent_scale` is taken from `cache/train/meta.json` (must match training).
+Only `latent_scale` is needed beyond the checkpoint weights/config; architecture
+comes from `ckpt["config"]`.
+
+Verify the Hub/export weights and settings against the original checkpoint:
+
+```bash
+python qadit_hf/verify_parity.py \
+  --ckpt path/to/ckpt_0023999.pt \
+  --model USER/qadit
+```
+
+Add `--forward --device cuda` for a full numerical DiT forward comparison.
 
 ## Files
 

config.json

@@ -1,38 +1,36 @@
 {
-  "architectures": [
-    "QaDiTModel"
-  ],
+  "architectures": ["QaDiTModel"],
+  "model_type": "qadit",
   "auto_map": {
     "AutoConfig": "configuration_qadit.QaDiTConfig",
     "AutoModel": "modelling_qadit.QaDiTModel"
   },
-  "depth": 12,
-  "dtype": "float32",
-  "duration_s": 10.24,
-  "guidance_scale": 4.0,
-  "hidden_size": 768,
   "latent_channels": 8,
-  "latent_freq": 16,
-  "latent_scale": 0.9503,
   "latent_time": 256,
-  "load_auxiliaries": true,
-  "logit_normal_mean": 0.0,
-  "logit_normal_std": 1.0,
-  "mlp_ratio": 4.0,
-  "model_type": "qadit",
-  "num_heads": 12,
-  "num_inference_steps": 50,
-  "num_train_timesteps": 1000,
+  "latent_freq": 16,
+  "latent_scale": 0.9503493000009796,
   "patch_size": 2,
+  "hidden_size": 768,
+  "depth": 12,
+  "num_heads": 12,
+  "mlp_ratio": 4.0,
   "repa_layer": 4,
-  "sample_rate": 16000,
-  "schedule": "cosine",
+  "text_model": "google/flan-t5-large",
   "text_dim": 1024,
   "text_max_length": 64,
-  "text_model": "google/flan-t5-large",
-  "transformers_version": "5.14.1",
   "vae_model": "cvssp/audioldm-s-full-v2",
   "vae_subfolder": "vae",
   "vocoder_model": "cvssp/audioldm-s-full-v2",
-  "vocoder_subfolder": "vocoder"
+  "vocoder_subfolder": "vocoder",
+  "sample_rate": 16000,
+  "duration_s": 10.24,
+  "num_train_timesteps": 1000,
+  "schedule": "cosine",
+  "logit_normal_mean": 0.0,
+  "logit_normal_std": 1.0,
+  "num_inference_steps": 50,
+  "guidance_scale": 4.0,
+  "load_auxiliaries": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0"
 }

configuration_qadit.py

@@ -19,7 +19,9 @@ class QaDiTConfig(PretrainedConfig):
         latent_channels: int = 8,
         latent_time: int = 256,
         latent_freq: int = 16,
-        latent_scale: float = 1.0,
+        # Exact scale embedded in the released step-23999 training checkpoint.
+        # Converted checkpoints override this from their own train metadata.
+        latent_scale: float = 0.9503493000009796,
         # DiT-B (~160M)
         patch_size: int = 2,
         hidden_size: int = 768,

modelling_qadit.py

@@ -656,6 +656,7 @@ class QaDiTModel(PreTrainedModel):
         encoder_attention_mask: Optional[torch.Tensor] = None,
         num_inference_steps: Optional[int] = None,
         guidance_scale: Optional[float] = None,
+        seed: Optional[int] = 0,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         eta: float = 0.0,
         output_type: str = "np",
@@ -673,16 +674,26 @@ class QaDiTModel(PreTrainedModel):
             DDIM steps (default from config).
         guidance_scale:
             Classifier-free guidance scale (default from config).
+        seed:
+            Random seed used when ``generator`` is not supplied. Defaults to
+            0, matching the original ``audio_dit/sample.py`` CLI.
         output_type:
             ``"np"`` → numpy waveforms, ``"pt"`` → torch waveforms,
             ``"latent"`` → scaled latents only (no VAE/vocoder).
         """
         cfg = self.config
         device = self.device
+        # from_pretrained() normally returns eval mode, but make generation
+        # invariant to callers having toggled train() in the same process.
+        self.eval()
         steps = num_inference_steps or cfg.num_inference_steps
         guidance = (
             guidance_scale if guidance_scale is not None else cfg.guidance_scale
         )
+        if cfg.latent_scale <= 0:
+            raise ValueError(
+                f"config.latent_scale must be positive, got {cfg.latent_scale}"
+            )
 
         if encoder_hidden_states is None:
             if prompt is None:
@@ -713,6 +724,8 @@ class QaDiTModel(PreTrainedModel):
             cfg.latent_freq,
         )
 
+        if generator is None and seed is not None:
+            generator = torch.Generator(device=device.type).manual_seed(seed)
         if isinstance(generator, list):
             if len(generator) != B:
                 raise ValueError(