fix(diffusion): resolve #1671 TCD clone + default-construction timeouts

src/Diffusion/FastGeneration/FluxSchnellModel.cs

@@ -58,6 +58,10 @@ public class FluxSchnellModel<T> : LatentDiffusionModelBase<T>
     private FluxDoubleStreamPredictor<T> _predictor;
     private StandardVAE<T> _vae;
     private readonly IConditioningModule<T>? _conditioner;
+    // Resolved (never-null) construction seed for the default predictor/VAE. Clone() reuses it so a
+    // never-forwarded — still lazy — predictor materializes IDENTICAL weights in the clone instead of
+    // diverging on a fresh seed. Resolved once, so two separately-constructed unseeded models still differ.
+    private readonly int _layerSeed;
 
     /// <inheritdoc />
     public override INoisePredictor<T> NoisePredictor => _predictor;
@@ -88,12 +92,14 @@ public FluxSchnellModel(
             architecture)
     {
         _conditioner = conditioner;
-        InitializeLayers(predictor, vae, seed);
+        // Resolve null → a concrete seed so Clone() can reconstruct identical lazy weights from it.
+        _layerSeed = seed ?? RandomGenerator.Next();
+        InitializeLayers(predictor, vae, _layerSeed);
         SetGuidanceScale(DEFAULT_GUIDANCE);
     }
 
     [MemberNotNull(nameof(_predictor), nameof(_vae))]
-    private void InitializeLayers(FluxDoubleStreamPredictor<T>? predictor, StandardVAE<T>? vae, int? seed)
+    private void InitializeLayers(FluxDoubleStreamPredictor<T>? predictor, StandardVAE<T>? vae, int seed)
     {
         _predictor = predictor ?? new FluxDoubleStreamPredictor<T>(
             variant: FluxPredictorVariant.Schnell,
@@ -139,11 +145,18 @@ public override void SetParameters(Vector<T> parameters)
     /// <inheritdoc />
     public override IDiffusionModel<T> Clone()
     {
-        var clone = new FluxSchnellModel<T>(conditioner: _conditioner, seed: RandomGenerator.Next());
-        // Field-by-field clone — bypasses the int-bounded flat
-        // Vector<T> that GetParameters/SetParameters round-trip would
-        // require for ~12B FLUX-scale weights.
-        clone._predictor.SetParameters(_predictor.GetParameters());
+        // Reuse THIS model's resolved construction seed (not a fresh one): a never-forwarded predictor is
+        // still lazy, so the clone must materialize from the SAME seed to stay equivalent — a fresh seed
+        // would later materialize different weights and break Clone() equivalence.
+        var clone = new FluxSchnellModel<T>(conditioner: _conditioner, seed: _layerSeed);
+        // Scale-safe + lazy-preserving: only copy the foundation-scale (~12B-param FLUX) predictor's
+        // weights if they were actually materialized. A never-forwarded model's weights are still lazy, so
+        // the clone reconstructs them from the shared seed above (nothing to copy) — copying would
+        // pointlessly materialize the predictor twice (source + clone) and OOM. When a copy IS needed it
+        // streams per-tensor chunks (#1624), never the int-bounded flat Vector<T> that
+        // SetParameters(GetParameters()) builds (which threw "Array dimensions exceeded supported range").
+        if (_predictor.WeightsMaterialized)
+            clone._predictor.SetParameterChunks(_predictor.GetParameterChunks());
         clone._vae.SetParameters(_vae.GetParameters());
         return clone;
     }

src/Diffusion/FastGeneration/TCDModel.cs

@@ -53,6 +53,18 @@ public class TCDModel<T> : LatentDiffusionModelBase<T>
     private const int LATENT_CHANNELS = 4;
     private const double DEFAULT_GUIDANCE = 0.0;
 
+    /// <summary>
+    /// Paper-optimal sampling step count. TCD is a trajectory-consistency
+    /// distillation model designed for high-quality few-step generation; the
+    /// authors report 4 steps as the sweet spot (8 max recommended). The base
+    /// <see cref="DiffusionModelOptions{T}.DefaultInferenceSteps"/> is the
+    /// generic 10-step DDIM default, which both wastes ~2.5× the compute TCD
+    /// needs and is not how the distilled model is meant to be sampled — so the
+    /// ctor pins this value, matching every other few-step model in this folder
+    /// (e.g. ConsistencyModel = 2, FluxSchnell = 4).
+    /// </summary>
+    private const int OPTIMAL_INFERENCE_STEPS = 4;
+
     private UNetNoisePredictor<T> _predictor;
     private StandardVAE<T> _vae;
     private readonly IConditioningModule<T>? _conditioner;
@@ -83,7 +95,8 @@ public TCDModel(
             options ?? new DiffusionModelOptions<T>
             {
                 TrainTimesteps = 1000, BetaStart = 0.00085,
-                BetaEnd = 0.012, BetaSchedule = BetaSchedule.ScaledLinear
+                BetaEnd = 0.012, BetaSchedule = BetaSchedule.ScaledLinear,
+                DefaultInferenceSteps = OPTIMAL_INFERENCE_STEPS
             },
             scheduler ?? new DDIMScheduler<T>(SchedulerConfig<T>.CreateStableDiffusion()),
             architecture)
@@ -159,7 +172,7 @@ public override ModelMetadata<T> GetModelMetadata()
         m.SetProperty("text_encoder", "CLIP ViT-L/14");
         m.SetProperty("context_dim", 768);
         m.SetProperty("distillation_method", "trajectory-consistency");
-        m.SetProperty("optimal_steps", 4);
+        m.SetProperty("optimal_steps", OPTIMAL_INFERENCE_STEPS);
         m.SetProperty("max_recommended_steps", 8);
         m.SetProperty("guidance_scale", DEFAULT_GUIDANCE);
         m.SetProperty("latent_channels", LATENT_CHANNELS);

src/Diffusion/NoisePredictors/FlagDiTPredictor.cs

@@ -156,7 +156,13 @@ private void InitializeLayers(int? seed)
         {
             _attnNormPre[i] = new RMSNormalizationLayer<T>(_hiddenSize);
             _attnNormPost[i] = new RMSNormalizationLayer<T>(_hiddenSize);
-            _attn[i] = new GroupedQueryAttentionLayer<T>(_seqLen, _hiddenSize, _numHeads, _numKVHeads);
+            // deferAllocation: keep the GQA projection weights zero-sized until the first forward,
+            // matching the LazyDense FFN/adaLN layers above. Without this the 32-layer × 4096-hidden
+            // Flag-DiT stack eagerly allocates ~1.3 B attention weights in the constructor (#1671:
+            // DefaultConstruction >10 s timeout); the weight-streaming forward path materializes them
+            // on demand.
+            _attn[i] = new GroupedQueryAttentionLayer<T>(_seqLen, _hiddenSize, _numHeads, _numKVHeads,
+                deferAllocation: true);
             // RoPE: resolution-agnostic rotary position embedding (Su et al. 2021; Flag-DiT §3.1).
             _attn[i].ConfigurePositionalEncoding(PositionalEncodingType.Rotary, ropeTheta: 10000.0,
                 maxSequenceLength: System.Math.Max(_seqLen, 64));

src/Diffusion/NoisePredictors/MMDiTNoisePredictor.cs

@@ -1195,16 +1195,28 @@ private IEnumerable<ILayer<T>> MMDiTLayerSequence()
         yield return _outputProj;
     }
 
+    /// <summary>
+    /// True once this predictor's lazy weights have been materialized (its patch-embed is initialized,
+    /// which the first forward triggers). A never-materialized foundation-scale model has nothing to copy,
+    /// so internal callers (e.g. a wrapping model's <c>Clone</c>) can skip the multi-GB parameter copy and
+    /// stay lazy. Internal: this is clone/streaming materialization plumbing, not public model behavior.
+    /// </summary>
+    internal bool WeightsMaterialized => _patchEmbed.IsInitialized;
+
     /// <inheritdoc />
     public override IEnumerable<Tensor<T>> GetParameterChunks()
     {
-        // #1624: one chunk per layer in the canonical MMDiTLayerSequence order, so the flat
-        // concatenation is index-identical to GetParameters without materializing the full
-        // multi-billion-parameter aggregate that overflows/OOMs at default size.
+        // #1624 zero-copy: materialize each layer's lazy weights, then yield its resident trainable
+        // tensors BY REFERENCE — one chunk per tensor, in canonical MMDiTLayerSequence × GetTrainable
+        // order — instead of concatenating each layer's params into a transient multi-GB Vector<T>
+        // (which GC-thrashes/OOMs at >2.1B FLUX/MMDiT scale). Consumers (Clone, LatentDiffusionModelBase)
+        // pair Get/SetParameterChunks and count chunks dynamically, so per-tensor framing is consistent.
         foreach (var layer in MMDiTLayerSequence())
         {
-            var p = layer.GetParameters();
-            if (p.Length > 0) yield return new Tensor<T>(new[] { p.Length }, p);
+            if (layer is not LayerBase<T> lb) continue;
+            lb.MaterializeParameters();
+            foreach (var t in lb.GetTrainableParameters())
+                if (t.Length > 0) yield return t;
         }
     }
 
@@ -1214,16 +1226,30 @@ public override void SetParameterChunks(IEnumerable<Tensor<T>> chunks)
         using var e = chunks.GetEnumerator();
         foreach (var layer in MMDiTLayerSequence())
         {
-            if (layer.ParameterCount == 0) continue;
-            if (!e.MoveNext())
-                throw new System.ArgumentException(
-                    "SetParameterChunks received fewer chunks than MMDiT has parameterized layers.",
-                    nameof(chunks));
-            layer.SetParameters(e.Current.ToVector());
+            if (layer is not LayerBase<T> lb) continue;
+            lb.MaterializeParameters();
+            var dst = lb.GetTrainableParameters();
+            // Pull one chunk per non-empty trainable tensor and copy the values IN PLACE
+            // (CopyTrainableParametersFrom: no rebinding — which would alias clone↔source — and no flat
+            // aggregate). Empty slots stay aligned to keep the per-tensor index identical to the getter.
+            bool anyNonEmpty = false;
+            foreach (var t in dst) if (t.Length > 0) { anyNonEmpty = true; break; }
+            if (!anyNonEmpty) continue;
+            var incoming = new Tensor<T>[dst.Count];
+            for (int i = 0; i < dst.Count; i++)
+            {
+                if (dst[i].Length == 0) { incoming[i] = dst[i]; continue; }
+                if (!e.MoveNext())
+                    throw new System.ArgumentException(
+                        "SetParameterChunks received fewer chunks than MMDiT has parameter tensors.",
+                        nameof(chunks));
+                incoming[i] = e.Current;
+            }
+            lb.CopyTrainableParametersFrom(incoming);
         }
         if (e.MoveNext())
             throw new System.ArgumentException(
-                "SetParameterChunks received more chunks than MMDiT has parameterized layers.",
+                "SetParameterChunks received more chunks than MMDiT has parameter tensors.",
                 nameof(chunks));
     }