(self, unet=None, vae=None, clip=None, clip_vision=None, filename=None, vae_filename=None)
| 36 | |
| 37 | class StableDiffusionModel: |
| 38 | def __init__(self, unet=None, vae=None, clip=None, clip_vision=None, filename=None, vae_filename=None): |
| 39 | self.unet = unet |
| 40 | self.vae = vae |
| 41 | self.clip = clip |
| 42 | self.clip_vision = clip_vision |
| 43 | self.filename = filename |
| 44 | self.vae_filename = vae_filename |
| 45 | self.unet_with_lora = unet |
| 46 | self.clip_with_lora = clip |
| 47 | self.visited_loras = '' |
| 48 | |
| 49 | self.lora_key_map_unet = {} |
| 50 | self.lora_key_map_clip = {} |
| 51 | |
| 52 | if self.unet is not None: |
| 53 | self.lora_key_map_unet = model_lora_keys_unet(self.unet.model, self.lora_key_map_unet) |
| 54 | self.lora_key_map_unet.update({x: x for x in self.unet.model.state_dict().keys()}) |
| 55 | |
| 56 | if self.clip is not None: |
| 57 | self.lora_key_map_clip = model_lora_keys_clip(self.clip.cond_stage_model, self.lora_key_map_clip) |
| 58 | self.lora_key_map_clip.update({x: x for x in self.clip.cond_stage_model.state_dict().keys()}) |
| 59 | |
| 60 | @torch.no_grad() |
| 61 | @torch.inference_mode() |
nothing calls this directly
no test coverage detected