Skip to content

Commit 41ea2f8

Browse files
committed
update ESRGAN
1 parent ec352cf commit 41ea2f8

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

diffsynth/extensions/ESRGAN/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,12 @@ def decode_images(self, images):
107107

108108
@torch.no_grad()
109109
def upscale(self, images, batch_size=4, progress_bar=lambda x:x):
110+
if not isinstance(images, list):
111+
images = [images]
112+
is_single_image = True
113+
else:
114+
is_single_image = False
115+
110116
# Preprocess
111117
input_tensor = self.process_images(images)
112118

@@ -126,4 +132,6 @@ def upscale(self, images, batch_size=4, progress_bar=lambda x:x):
126132

127133
# To images
128134
output_images = self.decode_images(output_tensor)
135+
if is_single_image:
136+
output_images = output_images[0]
129137
return output_images

0 commit comments

Comments
 (0)