| 244 | return os.path.join(head, new_file_name) |
| 245 | |
| 246 | class InstructPix2Pix: |
| 247 | def __init__(self, device): |
| 248 | print(f"Initializing InstructPix2Pix to {device}") |
| 249 | self.device = device |
| 250 | self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
| 251 | |
| 252 | self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", |
| 253 | safety_checker=StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker'), |
| 254 | torch_dtype=self.torch_dtype).to(device) |
| 255 | self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) |
| 256 | |
| 257 | @prompts(name="Instruct Image Using Text", |
| 258 | description="useful when you want to the style of the image to be like the text. " |
| 259 | "like: make it look like a painting. or make it like a robot. " |
| 260 | "The input to this tool should be a comma separated string of two, " |
| 261 | "representing the image_path and the text. ") |
| 262 | def inference(self, inputs): |
| 263 | """Change style of image.""" |
| 264 | print("===>Starting InstructPix2Pix Inference") |
| 265 | image_path, text = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) |
| 266 | original_image = Image.open(image_path) |
| 267 | image = self.pipe(text, image=original_image, num_inference_steps=40, image_guidance_scale=1.2).images[0] |
| 268 | updated_image_path = get_new_image_name(image_path, func_name="pix2pix") |
| 269 | image.save(updated_image_path) |
| 270 | print(f"\nProcessed InstructPix2Pix, Input Image: {image_path}, Instruct Text: {text}, " |
| 271 | f"Output Image: {updated_image_path}") |
| 272 | return updated_image_path |
| 273 | |
| 274 | |
| 275 | class Text2Image: |
nothing calls this directly
no outgoing calls
no test coverage detected