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)