|
from modules.utils import * |
|
|
|
class InstructPix2Pix: |
|
def __init__(self, device, pretrained_model_dir): |
|
print("Initializing InstructPix2Pix to %s" % device) |
|
self.device = device |
|
self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
|
self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(f"{pretrained_model_dir}/instruct-pix2pix", |
|
safety_checker=None, |
|
torch_dtype=self.torch_dtype).to(device) |
|
self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) |
|
|
|
@prompts(name="Instruct Image Using Text", |
|
description="useful when you want to the style of the image to be like the text. " |
|
"like: make it look like a painting. or make it like a robot. " |
|
"The input to this tool should be a comma seperated string of two, " |
|
"representing the image_path and the text. ") |
|
def inference(self, inputs): |
|
"""Change style of image.""" |
|
print("===>Starting InstructPix2Pix Inference") |
|
image_path, text = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) |
|
original_image = Image.open(image_path) |
|
image = self.pipe(text, image=original_image, num_inference_steps=40, image_guidance_scale=1.2).images[0] |
|
updated_image_path = get_new_image_name(image_path, func_name="pix2pix") |
|
image.save(updated_image_path) |
|
print(f"\nProcessed InstructPix2Pix, Input Image: {image_path}, Instruct Text: {text}, " |
|
f"Output Image: {updated_image_path}") |
|
return updated_image_path |
