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" => {
29            return Some("flow.out_adaln.weight".into())
30        }
31        "flow_loss.net.final_layer.adaLN_modulation.1.bias" => {
32            return Some("flow.out_adaln.bias".into())
33        }
34        _ => {}
35    }
36
37    // Transformer blocks
38    if let Some(rest) = hf_name.strip_prefix("model.layers.") {
39        if let Some((n_str, rest)) = rest.split_once('.') {
40            if let Ok(n) = n_str.parse::<u32>() {
41                let gguf = match rest {
42                    "self_attn.q_proj.weight" => "attn_q.weight",
43                    "self_attn.q_proj.bias" => "attn_q.bias",
44                    "self_attn.k_proj.weight" => "attn_k.weight",
45                    "self_attn.k_proj.bias" => "attn_k.bias",
46                    "self_attn.v_proj.weight" => "attn_v.weight",
47                    "self_attn.v_proj.bias" => "attn_v.bias",
48                    "self_attn.o_proj.weight" => "attn_out.weight",
49                    "norm1.weight" => "attn_norm.weight",
50                    "norm1.bias" => "attn_norm.bias",
51                    "norm2.weight" => "ffn_norm.weight",
52                    "norm2.bias" => "ffn_norm.bias",
53                    "ffn_layer.gate_proj.weight" => "ffn_gate.weight",
54                    "ffn_layer.up_proj.weight" => "ffn_up.weight",
55                    "ffn_layer.down_proj.weight" => "ffn_down.weight",
56                    _ => return None,
57                };
58                return Some(format!("blk.{n}.{gguf}"));
59            }
60        }
61    }
62
63    // Flow residual blocks
64    if let Some(rest) = hf_name.strip_prefix("flow_loss.net.res_blocks.") {
65        if let Some((k_str, rest)) = rest.split_once('.') {
66            if let Ok(k) = k_str.parse::<u32>() {
67                let gguf = match rest {
68                    "in_ln.weight" => "ln.weight",
69                    "in_ln.bias" => "ln.bias",
70                    "mlp.0.weight" => "mlp1.weight",
71                    "mlp.0.bias" => "mlp1.bias",
72                    "mlp.2.weight" => "mlp2.weight",
73                    "mlp.2.bias" => "mlp2.bias",
74                    "adaLN_modulation.1.weight" => "adaln.weight",
75                    "adaLN_modulation.1.bias" => "adaln.bias",
76                    _ => return None,
77                };
78                return Some(format!("flow.res.{k}.{gguf}"));
79            }
80        }
81    }
82
83    None
84}