From 896b0f684b41cc7a7638626ff298b46eb2298221 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Mon, 19 Dec 2022 12:42:31 -0300 Subject: [PATCH] Minor fixes to upscaler cli tool --- skynet/dgpu.py | 2 +- skynet/utils.py | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/skynet/dgpu.py b/skynet/dgpu.py index a26336e..7b44d6c 100644 --- a/skynet/dgpu.py +++ b/skynet/dgpu.py @@ -56,7 +56,7 @@ def pipeline_for(algo: str, mem_fraction: float = 1.0): pipe.enable_vae_slicing() - return pipe.to("cuda") + return pipe.to('cuda') class DGPUComputeError(BaseException): diff --git a/skynet/utils.py b/skynet/utils.py index 0534160..06bba1d 100644 --- a/skynet/utils.py +++ b/skynet/utils.py @@ -64,18 +64,19 @@ def upscale( torch.backends.cudnn.allow_tf32 = True login(token=hf_token) - params = { - 'torch_dtype': torch.float16, - 'safety_checker': None - } pipe = StableDiffusionUpscalePipeline.from_pretrained( - 'stabilityai/stable-diffusion-x4-upscaler', **params) + 'stabilityai/stable-diffusion-x4-upscaler', + revision="fp16", torch_dtype=torch.float16) + pipe.scheduler = EulerAncestralDiscreteScheduler.from_config( + pipe.scheduler.config) + pipe = pipe.to('cuda') prompt = prompt image = pipe( prompt, - image=Image.open(img_path) + image=Image.open(img_path).convert("RGB"), + num_inference_steps=steps ).images[0] image.save(output)