| 73 | keys_ignore = ["_extra_state", "accum_"] |
| 74 | |
| 75 | class ModelHyVid(ModelTemplate): |
| 76 | arch = "hyvid" |
| 77 | keys_detect = [ |
| 78 | ( |
| 79 | "double_blocks.0.img_attn_proj.weight", |
| 80 | "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", |
| 81 | ) |
| 82 | ] |
| 83 | |
| 84 | def handle_nd_tensor(self, key, data): |
| 85 | # hacky but don't have any better ideas |
| 86 | path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? |
| 87 | if os.path.isfile(path): |
| 88 | raise RuntimeError(f"5D tensor fix file already exists! {path}") |
| 89 | fsd = {key: torch.from_numpy(data)} |
| 90 | tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") |
| 91 | save_file(fsd, path) |
| 92 | |
| 93 | class ModelWan(ModelHyVid): |
| 94 | arch = "wan" |
nothing calls this directly
no outgoing calls
no test coverage detected