Merge branch 'dev' into extra-networks-always-visible

1 files changed, 13 insertions, 7 deletions

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 763829f1..b3d344e7 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -2,10 +2,8 @@ from collections import namedtuple
 import numpy as np

 import torch

 from PIL import Image

-from modules import devices, processing, images, sd_vae_approx, sd_samplers, sd_vae_taesd

-

+from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared

 from modules.shared import opts, state

-import modules.shared as shared

 

 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])

 

@@ -37,7 +35,7 @@ def single_sample_to_image(sample, approximation=None):
         x_sample = sample * 1.5

         x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()

     else:

-        x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5

+        x_sample = decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5

 

     x_sample = torch.clamp(x_sample, min=0.0, max=1.0)

     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

@@ -46,6 +44,12 @@ def single_sample_to_image(sample, approximation=None):
     return Image.fromarray(x_sample)

 

 

+def decode_first_stage(model, x):

+    x = model.decode_first_stage(x.to(devices.dtype_vae))

+

+    return x

+

+

 def sample_to_image(samples, index=0, approximation=None):

     return single_sample_to_image(samples[index], approximation)

 

@@ -85,11 +89,13 @@ class InterruptedException(BaseException):
     pass

 

 

-if opts.randn_source == "CPU":

+def replace_torchsde_browinan():

     import torchsde._brownian.brownian_interval

 

     def torchsde_randn(size, dtype, device, seed):

-        generator = torch.Generator(devices.cpu).manual_seed(int(seed))

-        return torch.randn(size, dtype=dtype, device=devices.cpu, generator=generator).to(device)

+        return devices.randn_local(seed, size).to(device=device, dtype=dtype)

 

     torchsde._brownian.brownian_interval._randn = torchsde_randn

+

+

+replace_torchsde_browinan()