theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-23 17:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

update ESRGAN

Browse Source
This commit is contained in:
Artiprocher
2024-10-08 18:23:39 +08:00
parent ec352cfce2
commit 41ea2f811a
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
diffsynth/extensions/ESRGAN/__init__.py
View File

@@ -107,6 +107,12 @@ class ESRGAN(torch.nn.Module):
@torch.no_grad() @torch.no_grad()
def upscale(self, images, batch_size=4, progress_bar=lambda x:x): def upscale(self, images, batch_size=4, progress_bar=lambda x:x):
if not isinstance(images, list):
images = [images]
is_single_image = True
else:
is_single_image = False
# Preprocess # Preprocess
input_tensor = self.process_images(images) input_tensor = self.process_images(images)
@@ -126,4 +132,6 @@ class ESRGAN(torch.nn.Module):
# To images # To images
output_images = self.decode_images(output_tensor) output_images = self.decode_images(output_tensor)
if is_single_image:
output_images = output_images[0]
return output_images return output_images