Spaces:
Running
Running
File size: 880 Bytes
111afa2 |
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 |
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
|