diff --git a/ldm/invoke/generator/base.py b/ldm/invoke/generator/base.py index e326bcfe8d..2e96c93cbb 100644 --- a/ldm/invoke/generator/base.py +++ b/ldm/invoke/generator/base.py @@ -6,6 +6,7 @@ import torch import numpy as np import random import os +import traceback from tqdm import tqdm, trange from PIL import Image, ImageFilter from einops import rearrange, repeat @@ -82,7 +83,9 @@ class Generator(): try: x_T = self.get_noise(width,height) except: - pass + print('** An error occurred while getting initial noise **') + print(traceback.format_exc()) + image = make_image(x_T) if self.safety_checker is not None: diff --git a/ldm/invoke/generator/img2img.py b/ldm/invoke/generator/img2img.py index 31c3ca256e..73eb2e6a06 100644 --- a/ldm/invoke/generator/img2img.py +++ b/ldm/invoke/generator/img2img.py @@ -14,7 +14,7 @@ from ldm.models.diffusion.ddim import DDIMSampler class Img2Img(Generator): def __init__(self, model, precision): super().__init__(model, precision) - self.init_latent = None # by get_noise() + self.init_latent = None # by get_noise() def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta, conditioning,init_image,strength,step_callback=None,threshold=0.0,perlin=0.0,**kwargs): diff --git a/ldm/invoke/generator/omnibus.py b/ldm/invoke/generator/omnibus.py index d841cac64f..e0705ec397 100644 --- a/ldm/invoke/generator/omnibus.py +++ b/ldm/invoke/generator/omnibus.py @@ -70,6 +70,7 @@ class Omnibus(Img2Img,Txt2Img): mask_image = torch.ones(1, 1, height, width, device=self.model.device) masked_image = init_image + self.init_latent = init_image height = init_image.shape[2] width = init_image.shape[3] model = self.model @@ -144,4 +145,7 @@ class Omnibus(Img2Img,Txt2Img): return batch def get_noise(self, width:int, height:int): - return super(Txt2Img,self).get_noise(width,height) + if self.init_latent is not None: + height = self.init_latent.shape[2] + width = self.init_latent.shape[3] + return Txt2Img.get_noise(self,width,height)