[Tests] Config driven tiny model builder (#5090)

Signed-off-by: Nick Cao <ncao@redhat.com>
2026-07-22 17:11:53 +00:00
1 parent cc1939120f
commit f10320f53f
@@ -0,0 +1,25 @@
+{
+  "_class_name": "Flux2KleinPipeline",
+  "_diffusers_version": "0.37.0.dev0",
+  "is_distilled": true,
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3ForCausalLM"
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Flux2Transformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}
@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 16,
+  "hidden_act": "silu",
+  "hidden_size": 32,
+  "initializer_range": 0.02,
+  "intermediate_size": 64,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 512,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 2,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.56.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}
@@ -0,0 +1,24 @@
+{
+  "_class_name": "Flux2Transformer2DModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "_name_or_path": "/raid/yiyi/klein-4b-distilled-diffusers/transformer",
+  "attention_head_dim": 16,
+  "axes_dims_rope": [
+    4,
+    4,
+    4,
+    4
+  ],
+  "eps": 1e-06,
+  "guidance_embeds": false,
+  "in_channels": 64,
+  "joint_attention_dim": 96,
+  "mlp_ratio": 3.0,
+  "num_attention_heads": 2,
+  "num_layers": 2,
+  "num_single_layers": 2,
+  "out_channels": null,
+  "patch_size": 1,
+  "rope_theta": 2000,
+  "timestep_guidance_channels": 32
+}
@@ -0,0 +1,31 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.37.0.dev0",
+  "_name_or_path": "black-forest-labs/FLUX.2-dev",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    32
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "layers_per_block": 1,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 16,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}