File size: 1,976 Bytes
7fbdac4 c95142c 7fbdac4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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,
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}
|