Fix for crashes in txt2img hires fix mode

This commit is contained in:
ArDiouscuros 2022-10-05 18:31:04 +02:00 committed by Lincoln Stein
parent 0f55d89e20
commit 5a40f7ad15

View File

@ -40,7 +40,12 @@ class Txt2Img2Img(Generator):
init_width // self.downsampling_factor, init_width // self.downsampling_factor,
] ]
x = self.get_noise(init_width, init_height) sampler.make_schedule(
ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
)
#x = self.get_noise(init_width, init_height)
x = x_T
if self.free_gpu_mem and self.model.model.device != self.model.device: if self.free_gpu_mem and self.model.model.device != self.model.device:
self.model.model.to(self.model.device) self.model.model.to(self.model.device)
@ -71,7 +76,7 @@ class Txt2Img2Img(Generator):
t_enc = int(strength * steps) t_enc = int(strength * steps)
x = None x = self.get_noise(width,height,False)
# Other samplers not supported yet, so ignore previous sampler # Other samplers not supported yet, so ignore previous sampler
if not isinstance(sampler,DDIMSampler): if not isinstance(sampler,DDIMSampler):
@ -88,7 +93,7 @@ class Txt2Img2Img(Generator):
z_enc = img_sampler.stochastic_encode( z_enc = img_sampler.stochastic_encode(
samples, samples,
torch.tensor([t_enc]).to(self.model.device), torch.tensor([t_enc]).to(self.model.device),
noise=x_T noise=x
) )
# decode it # decode it
@ -110,17 +115,28 @@ class Txt2Img2Img(Generator):
# returns a tensor filled with random numbers from a normal distribution # returns a tensor filled with random numbers from a normal distribution
def get_noise(self,width,height): def get_noise(self,width,height,scale = True):
# print(f"Get noise: {width}x{height}")
if scale:
trained_square = 512 * 512
actual_square = width * height
scale = math.sqrt(trained_square / actual_square)
scaled_width = math.ceil(scale * width / 64) * 64
scaled_height = math.ceil(scale * height / 64) * 64
else:
scaled_width = width
scaled_height = height
device = self.model.device device = self.model.device
if device.type == 'mps': if device.type == 'mps':
return torch.randn([1, return torch.randn([1,
self.latent_channels, self.latent_channels,
height // self.downsampling_factor, scaled_height // self.downsampling_factor,
width // self.downsampling_factor], scaled_width // self.downsampling_factor],
device='cpu').to(device) device='cpu').to(device)
else: else:
return torch.randn([1, return torch.randn([1,
self.latent_channels, self.latent_channels,
height // self.downsampling_factor, scaled_height // self.downsampling_factor,
width // self.downsampling_factor], scaled_width // self.downsampling_factor],
device=device) device=device)