Spaces:
Runtime error
Runtime error
File size: 3,189 Bytes
c19ca42 |
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 65 66 67 68 69 70 71 72 73 74 |
import torch
import diffusers
from PIL import Image
from modules import shared, devices
from modules.upscaler import Upscaler, UpscalerData
class UpscalerSD(Upscaler):
def __init__(self, dirname): # pylint: disable=super-init-not-called
self.name = "SDUpscale"
self.user_path = dirname
if shared.backend != shared.Backend.DIFFUSERS:
super().__init__()
return
self.scalers = [
UpscalerData(name="SD Latent 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4),
UpscalerData(name="SD Latent 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4),
]
self.pipelines = [
None,
None,
]
self.models = {}
def load_model(self, path: str):
from modules.sd_models import set_diffuser_options
scaler: UpscalerData = [x for x in self.scalers if x.data_path == path][0]
if self.models.get(path, None) is not None:
shared.log.debug(f"Upscaler cached: type={scaler.name} model={path}")
return self.models[path]
else:
devices.set_cuda_params()
model = diffusers.DiffusionPipeline.from_pretrained(path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype)
if hasattr(model, "set_progress_bar_config"):
model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + 'Upscale', ncols=80, colour='#327fba')
set_diffuser_options(scaler.model, vae=None, op='upscaler')
self.models[path] = model
return self.models[path]
def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor):
pass
def do_upscale(self, img: Image.Image, selected_model):
devices.torch_gc()
model = self.load_model(selected_model)
if model is None:
return img
seeds = [torch.randint(0, 2 ** 32, (1,)).item() for _ in range(1)]
generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else devices.device
generator = [torch.Generator(generator_device).manual_seed(s) for s in seeds]
args = {
'prompt': '',
'negative_prompt': '',
'image': img,
'num_inference_steps': 20,
'guidance_scale': 7.5,
'generator': generator,
'latents': None,
'return_dict': True,
'callback': self.callback,
'callback_steps': 1,
# 'noise_level': 100,
# 'num_images_per_prompt': 1,
# 'eta': 0.0,
# 'cross_attention_kwargs': None,
}
model = model.to(devices.device)
output = model(**args)
image = output.images[0]
if shared.opts.upscaler_unload and selected_model in self.models:
del self.models[selected_model]
shared.log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}")
devices.torch_gc(force=True)
return image
|