Skip to main content

zsfm_sundial/
tensor_map.rs

1/// Map a HuggingFace Sundial tensor name to its GGUF equivalent.
2/// Returns None for tensors that should be skipped.
3pub fn map_tensor_name(hf_name: &str) -> Option<String> {
4    // Patch embedding (tokenizer)
5    match hf_name {
6        "model.embed_layer.hidden_layer.weight"   => return Some("embed.hidden.weight".into()),
7        "model.embed_layer.hidden_layer.bias"     => return Some("embed.hidden.bias".into()),
8        "model.embed_layer.output_layer.weight"   => return Some("embed.output.weight".into()),
9        "model.embed_layer.output_layer.bias"     => return Some("embed.output.bias".into()),
10        "model.embed_layer.residual_layer.weight" => return Some("embed.skip.weight".into()),
11        "model.embed_layer.residual_layer.bias"   => return Some("embed.skip.bias".into()),
12        // Final backbone norm
13        "model.norm.weight" => return Some("norm.weight".into()),
14        "model.norm.bias"   => return Some("norm.bias".into()),
15        // Flow head — timestep embedder
16        "flow_loss.net.time_embed.mlp.0.weight" => return Some("flow.t_proj1.weight".into()),
17        "flow_loss.net.time_embed.mlp.0.bias"   => return Some("flow.t_proj1.bias".into()),
18        "flow_loss.net.time_embed.mlp.2.weight" => return Some("flow.t_proj2.weight".into()),
19        "flow_loss.net.time_embed.mlp.2.bias"   => return Some("flow.t_proj2.bias".into()),
20        // Flow head — condition/input projections
21        "flow_loss.net.cond_embed.weight" => return Some("flow.cond.weight".into()),
22        "flow_loss.net.cond_embed.bias"   => return Some("flow.cond.bias".into()),
23        "flow_loss.net.input_proj.weight" => return Some("flow.in_proj.weight".into()),
24        "flow_loss.net.input_proj.bias"   => return Some("flow.in_proj.bias".into()),
25        // Flow head — final layer
26        "flow_loss.net.final_layer.linear.weight"             => return Some("flow.out_linear.weight".into()),
27        "flow_loss.net.final_layer.linear.bias"               => return Some("flow.out_linear.bias".into()),
28        "flow_loss.net.final_layer.adaLN_modulation.1.weight" => return Some("flow.out_adaln.weight".into()),
29        "flow_loss.net.final_layer.adaLN_modulation.1.bias"   => return Some("flow.out_adaln.bias".into()),
30        _ => {}
31    }
32
33    // Transformer blocks
34    if let Some(rest) = hf_name.strip_prefix("model.layers.") {
35        if let Some((n_str, rest)) = rest.split_once('.') {
36            if let Ok(n) = n_str.parse::<u32>() {
37                let gguf = match rest {
38                    "self_attn.q_proj.weight" => "attn_q.weight",
39                    "self_attn.q_proj.bias"   => "attn_q.bias",
40                    "self_attn.k_proj.weight" => "attn_k.weight",
41                    "self_attn.k_proj.bias"   => "attn_k.bias",
42                    "self_attn.v_proj.weight" => "attn_v.weight",
43                    "self_attn.v_proj.bias"   => "attn_v.bias",
44                    "self_attn.o_proj.weight" => "attn_out.weight",
45                    "norm1.weight"            => "attn_norm.weight",
46                    "norm1.bias"              => "attn_norm.bias",
47                    "norm2.weight"            => "ffn_norm.weight",
48                    "norm2.bias"              => "ffn_norm.bias",
49                    "ffn_layer.gate_proj.weight" => "ffn_gate.weight",
50                    "ffn_layer.up_proj.weight"   => "ffn_up.weight",
51                    "ffn_layer.down_proj.weight" => "ffn_down.weight",
52                    _ => return None,
53                };
54                return Some(format!("blk.{n}.{gguf}"));
55            }
56        }
57    }
58
59    // Flow residual blocks
60    if let Some(rest) = hf_name.strip_prefix("flow_loss.net.res_blocks.") {
61        if let Some((k_str, rest)) = rest.split_once('.') {
62            if let Ok(k) = k_str.parse::<u32>() {
63                let gguf = match rest {
64                    "in_ln.weight"              => "ln.weight",
65                    "in_ln.bias"                => "ln.bias",
66                    "mlp.0.weight"              => "mlp1.weight",
67                    "mlp.0.bias"                => "mlp1.bias",
68                    "mlp.2.weight"              => "mlp2.weight",
69                    "mlp.2.bias"                => "mlp2.bias",
70                    "adaLN_modulation.1.weight" => "adaln.weight",
71                    "adaLN_modulation.1.bias"   => "adaln.bias",
72                    _ => return None,
73                };
74                return Some(format!("flow.res.{k}.{gguf}"));
75            }
76        }
77    }
78
79    None
80}