|
import torch |
|
from diffusers import StableDiffusionInstructPix2PixPipeline |
|
|
|
import internals.util.image as ImageUtil |
|
from internals.data.dataAccessor import update_db |
|
from internals.data.task import Task |
|
from internals.util.cache import clear_cuda_and_gc |
|
from internals.util.commons import download_image, upload_images |
|
from internals.util.config import get_hf_token |
|
from internals.util.slack import Slack |
|
|
|
slack = Slack() |
|
|
|
|
|
class Script: |
|
def __init__(self, **kwargs): |
|
self.__name__ = "day_night_ip2p" |
|
|
|
@update_db |
|
@slack.auto_send_alert |
|
def __call__(self, task: Task, args: dict): |
|
clear_cuda_and_gc() |
|
|
|
model_id = args.get("model_id", None) |
|
steps = args.get("steps", 50) |
|
image_guidance_scale = args.get("image_guidance_scale", 1.5) |
|
guidance_scale = args.get("guidance_scale", 7.5) |
|
|
|
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained( |
|
model_id, |
|
use_auth_token=get_hf_token(), |
|
torch_dtype=torch.float16, |
|
safety_checker=None, |
|
).to("cuda") |
|
pipe.enable_xformers_memory_efficient_attention() |
|
|
|
prompt = ["convert to night", "convert to evening", "convert to midnight"] |
|
image = download_image(task.get_imageUrl()) |
|
image = ImageUtil.resize_image(image, 1024) |
|
|
|
images = [] |
|
for p in prompt: |
|
print("Generating: ", p) |
|
image = pipe.__call__( |
|
prompt=p, |
|
num_inference_steps=steps, |
|
image=image, |
|
guidance_scale=guidance_scale, |
|
num_images_per_prompt=1, |
|
image_guidance_scale=image_guidance_scale, |
|
).images[0] |
|
images.append(image) |
|
|
|
generated_image_urls = upload_images( |
|
images, "_" + self.__name__, task.get_taskId() |
|
) |
|
|
|
pipe = None |
|
del pipe |
|
|
|
clear_cuda_and_gc() |
|
|
|
return {"generated_image_urls": generated_image_urls} |
|
|
