(args)
| 241 | |
| 242 | |
| 243 | def main(args): |
| 244 | original_ckpt = load_original_checkpoint(args) |
| 245 | |
| 246 | if args.transformer: |
| 247 | num_layers = 6 |
| 248 | num_single_layers = 27 |
| 249 | inner_dim = 3072 |
| 250 | mlp_ratio = 4.0 |
| 251 | |
| 252 | converted_transformer_state_dict = convert_ovis_image_transformer_checkpoint_to_diffusers( |
| 253 | original_ckpt, num_layers, num_single_layers, inner_dim, mlp_ratio=mlp_ratio |
| 254 | ) |
| 255 | transformer = OvisImageTransformer2DModel(in_channels=args.in_channels, out_channels=args.out_channels) |
| 256 | transformer.load_state_dict(converted_transformer_state_dict, strict=True) |
| 257 | |
| 258 | print("Saving Ovis-Image Transformer in Diffusers format.") |
| 259 | transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") |
| 260 | |
| 261 | |
| 262 | if __name__ == "__main__": |
no test coverage detected
searching dependent graphs…