| 384 | |
| 385 | |
| 386 | class SD3DiTStateDictConverter: |
| 387 | def __init__(self): |
| 388 | pass |
| 389 | |
| 390 | def infer_architecture(self, state_dict): |
| 391 | embed_dim = state_dict["blocks.0.ff_a.0.weight"].shape[1] |
| 392 | num_layers = 100 |
| 393 | while num_layers > 0 and f"blocks.{num_layers-1}.ff_a.0.bias" not in state_dict: |
| 394 | num_layers -= 1 |
| 395 | use_rms_norm = "blocks.0.attn.norm_q_a.weight" in state_dict |
| 396 | num_dual_blocks = 0 |
| 397 | while f"blocks.{num_dual_blocks}.attn2.a_to_out.bias" in state_dict: |
| 398 | num_dual_blocks += 1 |
| 399 | pos_embed_max_size = state_dict["pos_embedder.pos_embed"].shape[1] |
| 400 | return { |
| 401 | "embed_dim": embed_dim, |
| 402 | "num_layers": num_layers, |
| 403 | "use_rms_norm": use_rms_norm, |
| 404 | "num_dual_blocks": num_dual_blocks, |
| 405 | "pos_embed_max_size": pos_embed_max_size |
| 406 | } |
| 407 | |
| 408 | def from_diffusers(self, state_dict): |
| 409 | rename_dict = { |
| 410 | "context_embedder": "context_embedder", |
| 411 | "pos_embed.pos_embed": "pos_embedder.pos_embed", |
| 412 | "pos_embed.proj": "pos_embedder.proj", |
| 413 | "time_text_embed.timestep_embedder.linear_1": "time_embedder.timestep_embedder.0", |
| 414 | "time_text_embed.timestep_embedder.linear_2": "time_embedder.timestep_embedder.2", |
| 415 | "time_text_embed.text_embedder.linear_1": "pooled_text_embedder.0", |
| 416 | "time_text_embed.text_embedder.linear_2": "pooled_text_embedder.2", |
| 417 | "norm_out.linear": "norm_out.linear", |
| 418 | "proj_out": "proj_out", |
| 419 | |
| 420 | "norm1.linear": "norm1_a.linear", |
| 421 | "norm1_context.linear": "norm1_b.linear", |
| 422 | "attn.to_q": "attn.a_to_q", |
| 423 | "attn.to_k": "attn.a_to_k", |
| 424 | "attn.to_v": "attn.a_to_v", |
| 425 | "attn.to_out.0": "attn.a_to_out", |
| 426 | "attn.add_q_proj": "attn.b_to_q", |
| 427 | "attn.add_k_proj": "attn.b_to_k", |
| 428 | "attn.add_v_proj": "attn.b_to_v", |
| 429 | "attn.to_add_out": "attn.b_to_out", |
| 430 | "ff.net.0.proj": "ff_a.0", |
| 431 | "ff.net.2": "ff_a.2", |
| 432 | "ff_context.net.0.proj": "ff_b.0", |
| 433 | "ff_context.net.2": "ff_b.2", |
| 434 | |
| 435 | "attn.norm_q": "attn.norm_q_a", |
| 436 | "attn.norm_k": "attn.norm_k_a", |
| 437 | "attn.norm_added_q": "attn.norm_q_b", |
| 438 | "attn.norm_added_k": "attn.norm_k_b", |
| 439 | } |
| 440 | state_dict_ = {} |
| 441 | for name, param in state_dict.items(): |
| 442 | if name in rename_dict: |
| 443 | if name == "pos_embed.pos_embed": |
no outgoing calls
no test coverage detected