import modal
import torch
from diffusers import StableDiffusionUpscalePipeline

from .app import app
from .image import image


@app.cls(gpu="T4", image=image, scaledown_window=60 * 5)
class UpscalerModalApp:
    @modal.enter()
    def setup(self):
        model_id = "stabilityai/stable-diffusion-x4-upscaler"
        self.pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
        self.pipeline = self.pipeline.to("cuda")

    @modal.batched(max_batch_size=4, wait_ms=1000)
    def forward(self, low_res_imgs, prompts: list[str]):
        print(len(low_res_imgs))
        print(low_res_imgs)
        print(prompts)
        low_res_imgs = [img.resize((min(512, img.width), min(512, img.height))) for img in low_res_imgs]
        upscaled_images = self.pipeline(prompt=prompts, image=low_res_imgs).images
        return upscaled_images