(self, device)
| 1170 | |
| 1171 | class Inpainting: |
| 1172 | def __init__(self, device): |
| 1173 | self.device = device |
| 1174 | self.revision = 'fp16' if 'cuda' in self.device else None |
| 1175 | self.torch_dtype = torch.float16 if 'cuda' in self.device else torch.float32 |
| 1176 | |
| 1177 | self.inpaint = StableDiffusionInpaintPipeline.from_pretrained( |
| 1178 | "runwayml/stable-diffusion-inpainting", revision=self.revision, torch_dtype=self.torch_dtype,safety_checker=StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker')).to(device) |
| 1179 | def __call__(self, prompt, image, mask_image, height=512, width=512, num_inference_steps=50): |
| 1180 | update_image = self.inpaint(prompt=prompt, image=image.resize((width, height)), |
| 1181 | mask_image=mask_image.resize((width, height)), height=height, width=width, num_inference_steps=num_inference_steps).images[0] |
nothing calls this directly
no outgoing calls
no test coverage detected