| 74 | layers_modified = [] |
| 75 | |
| 76 | def getlayers(model, root_name=''): |
| 77 | for name, module in model.named_children(): |
| 78 | if module.__class__.__name__ == 'SpatialTransformer': |
| 79 | layers_modified.append(root_name + '.' + name + '.transformer_blocks.0.attn2.to_k') |
| 80 | layers_modified.append(root_name + '.' + name + '.transformer_blocks.0.attn2.to_v') |
| 81 | else: |
| 82 | if list(module.children()) == []: |
| 83 | layers.append(root_name + '.' + name) |
| 84 | else: |
| 85 | getlayers(module, root_name + '.' + name) |
| 86 | |
| 87 | getlayers(model.model.diffusion_model) |
| 88 | |