File size: 1,808 Bytes
8d81eef |
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 |
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import torch
from PIL import Image
from refiners.foundationals.latent_diffusion.stable_diffusion_1.multi_upscaler import (
MultiUpscaler,
UpscalerCheckpoints,
)
from esrgan_model import UpscalerESRGAN
@dataclass(kw_only=True)
class ESRGANUpscalerCheckpoints(UpscalerCheckpoints):
"""Extends the SD-1.5 MultiUpscaler checkpoints to hold an extra ESRGAN file."""
esrgan: Path
class ESRGANUpscaler(MultiUpscaler):
"""
Multi-stage image enhancer that:
1. Runs ESRGAN 4× super-resolution first (tiling to avoid VRAM overflow),
2. Passes the up-scaled image to Stable-Diffusion 1.5 MultiUpscaler for refinement.
"""
def __init__(
self,
checkpoints: ESRGANUpscalerCheckpoints,
device: torch.device,
dtype: torch.dtype,
) -> None:
super().__init__(checkpoints=checkpoints, device=device, dtype=dtype)
self.esrgan = UpscalerESRGAN(
checkpoints.esrgan, device=self.device, dtype=self.dtype
)
# ---- automatically called by HF when the model is moved to another device ----
def to(self, device: torch.device, dtype: torch.dtype):
self.esrgan.to(device=device, dtype=dtype)
self.sd = self.sd.to(device=device, dtype=dtype)
self.device = device
self.dtype = dtype
# ---- hook that runs *before* SD-1.5 up-scaling ----
def pre_upscale(
self,
image: Image.Image,
upscale_factor: float,
**_: Any,
) -> Image.Image:
# 4× ESRGAN first, then the SD-1.5 stage handles the residual upscale
image = self.esrgan.upscale_with_tiling(image)
return super().pre_upscale(image=image, upscale_factor=upscale_factor / 4)
|