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}