Merge branch 'dev' into auro-autolaunch

1 files changed, 45 insertions, 11 deletions

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 763829f1..39586b40 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'])

 

@@ -25,19 +23,29 @@ def setup_img2img_steps(p, steps=None):
 approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}

 

 

-def single_sample_to_image(sample, approximation=None):

+def samples_to_images_tensor(sample, approximation=None, model=None):

+    '''latents -> images [-1, 1]'''

     if approximation is None:

         approximation = approximation_indexes.get(opts.show_progress_type, 0)

 

     if approximation == 2:

-        x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5

+        x_sample = sd_vae_approx.cheap_approximation(sample)

     elif approximation == 1:

-        x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach() * 0.5 + 0.5

+        x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype)).detach()

     elif approximation == 3:

         x_sample = sample * 1.5

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

+        x_sample = sd_vae_taesd.decoder_model()(x_sample.to(devices.device, devices.dtype)).detach()

+        x_sample = x_sample * 2 - 1

     else:

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

+        if model is None:

+            model = shared.sd_model

+        x_sample = model.decode_first_stage(sample.to(model.first_stage_model.dtype))

+

+    return x_sample

+

+

+def single_sample_to_image(sample, approximation=None):

+    x_sample = samples_to_images_tensor(sample.unsqueeze(0), approximation)[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 +54,12 @@ def single_sample_to_image(sample, approximation=None):
     return Image.fromarray(x_sample)

 

 

+def decode_first_stage(model, x):

+    x = x.to(devices.dtype_vae)

+    approx_index = approximation_indexes.get(opts.sd_vae_decode_method, 0)

+    return samples_to_images_tensor(x, approx_index, model)

+

+

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

     return single_sample_to_image(samples[index], approximation)

 

@@ -54,6 +68,24 @@ def samples_to_image_grid(samples, approximation=None):
     return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])

 

 

+def images_tensor_to_samples(image, approximation=None, model=None):

+    '''image[0, 1] -> latent'''

+    if approximation is None:

+        approximation = approximation_indexes.get(opts.sd_vae_encode_method, 0)

+

+    if approximation == 3:

+        image = image.to(devices.device, devices.dtype)

+        x_latent = sd_vae_taesd.encoder_model()(image)

+    else:

+        if model is None:

+            model = shared.sd_model

+        image = image.to(shared.device, dtype=devices.dtype_vae)

+        image = image * 2 - 1

+        x_latent = model.get_first_stage_encoding(model.encode_first_stage(image))

+

+    return x_latent

+

+

 def store_latent(decoded):

     state.current_latent = decoded

 

@@ -85,11 +117,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()