chunzhi/stable-diffusion-cli-on-modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Rename some directories.

Browse Source
This commit is contained in:
hodanov
2023-12-29 13:06:11 +09:00
parent ada08bd2fe
commit 5269f74630
11 changed files with 0 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/Dockerfile
+14
View File
@@ -0,0 +1,14 @@
FROM python:3.11.3-slim-bullseye
COPY ./requirements.txt /
RUN apt-get update \
&& apt-get install wget libgl1-mesa-glx libglib2.0-0 --no-install-recommends -y \
&& apt-get autoremove -y \
&& apt-get clean -y \
&& rm -rf /var/lib/apt/lists/* \
&& pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 --no-cache-dir \
&& mkdir -p /vol/cache/esrgan \
&& wget --progress=dot:giga https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /vol/cache/esrgan \
&& wget --progress=dot:giga https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth -P /vol/cache/esrgan \
&& wget --progress=dot:giga https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P /vol/cache/esrgan \
&& wget --progress=dot:giga https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth -P /vol/cache/esrgan \
&& wget --progress=dot:giga https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P /vol/cache/esrgan
app/__main__.py
+15
View File
@@ -0,0 +1,15 @@
from __future__ import annotations
import stable_diffusion_1_5
import stable_diffusion_xl
from setup import stub
@stub.function(gpu="A10G")
def main():
stable_diffusion_1_5.SD15
stable_diffusion_xl.SDXLTxt2Img
if __name__ == "__main__":
main.local()
app/config.sample.yml
+34
View File
@@ -0,0 +1,34 @@
##########
# This is the config file to set a base model, vae and some tools.
# Rename the file to `config.yml` before running the script.
# Execute `modal deploy ./setup_files/setup.py` every time modify this file.
##########
##########
# You can use a diffusers model and VAE on hugging face.
model:
name: stable-diffusion-1-5
url: https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned.safetensors
vae:
name: sd-vae-ft-mse
url: https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors
##########
# Add LoRA if you want to use one. You can use a download url such as the below.
# ex)
# loras:
# - name: hogehoge.safetensors
# url: https://hogehoge/xxxx
# - name: fugafuga.safetensors
# url: https://fugafuga/xxxx
##########
# You can use Textual Inversion and ControlNet also. Usage is the same as `loras`.
# ex)
# textual_inversions:
# - name: hogehoge
# url: https://hogehoge/xxxx
# - name: fugafuga
# url: https://fugafuga/xxxx
controlnets:
- name: control_v11f1e_sd15_tile
repo_id: lllyasviel/control_v11f1e_sd15_tile
app/requirements.txt
+25
View File
@@ -0,0 +1,25 @@
invisible_watermark
accelerate
diffusers[torch]==0.24.0
onnxruntime==1.16.3
safetensors==0.4.1
torch==2.1.0
transformers==4.36.2
xformers==0.0.22.post7
realesrgan==0.3.0
basicsr>=1.4.2
facexlib>=0.3.0
gfpgan>=1.3.8
scipy==1.11.4
opencv-python
Pillow
pillow-avif-plugin
torchvision
tqdm
controlnet_aux
pyyaml
# Use the below in 'download_from_original_stable_diffusion_ckpt'.
omegaconf==2.3.0
app/setup.py
+148
View File
@@ -0,0 +1,148 @@
from __future__ import annotations
import os
import diffusers
from modal import Image, Mount, Secret, Stub
BASE_CACHE_PATH = "/vol/cache"
BASE_CACHE_PATH_LORA = "/vol/cache/lora"
BASE_CACHE_PATH_TEXTUAL_INVERSION = "/vol/cache/textual_inversion"
BASE_CACHE_PATH_CONTROLNET = "/vol/cache/controlnet"
def download_file(url, file_name, file_path):
"""
Download files.
"""
from urllib.request import Request, urlopen
req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
downloaded = urlopen(req).read()
dir_names = os.path.join(file_path, file_name)
os.makedirs(os.path.dirname(dir_names), exist_ok=True)
with open(dir_names, mode="wb") as f:
f.write(downloaded)
def download_controlnet(name: str, repo_id: str, token: str):
"""
Download a controlnet.
"""
cache_path = os.path.join(BASE_CACHE_PATH_CONTROLNET, name)
controlnet = diffusers.ControlNetModel.from_pretrained(
repo_id,
use_auth_token=token,
cache_dir=cache_path,
)
controlnet.save_pretrained(cache_path, safe_serialization=True)
def download_vae(name: str, model_url: str, token: str):
"""
Download a vae.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
vae = diffusers.AutoencoderKL.from_single_file(
pretrained_model_link_or_path=model_url,
use_auth_token=token,
cache_dir=cache_path,
)
vae.save_pretrained(cache_path, safe_serialization=True)
def download_model(name: str, model_url: str, token: str):
"""
Download a model.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
pipe = diffusers.StableDiffusionPipeline.from_single_file(
pretrained_model_link_or_path=model_url,
use_auth_token=token,
cache_dir=cache_path,
)
pipe.save_pretrained(cache_path, safe_serialization=True)
def download_model_sdxl(name: str, model_url: str, token: str):
"""
Download a sdxl model.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=model_url,
use_auth_token=token,
cache_dir=cache_path,
)
pipe.save_pretrained(cache_path, safe_serialization=True)
refiner_cache_path = cache_path + "-refiner"
refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_single_file(
"https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/sd_xl_refiner_1.0.safetensors",
cache_dir=refiner_cache_path,
)
refiner.save_pretrained(refiner_cache_path, safe_serialization=True)
def build_image():
"""
Build the Docker image.
"""
import yaml
token = os.environ["HUGGING_FACE_TOKEN"]
config = {}
with open("/config.yml", "r") as file:
config = yaml.safe_load(file)
model = config.get("model")
use_xl = config.get("use_xl")
if model is not None:
if use_xl is not None and use_xl:
download_model_sdxl(name=model["name"], model_url=model["url"], token=token)
else:
download_model(name=model["name"], model_url=model["url"], token=token)
vae = config.get("vae")
if vae is not None:
download_vae(name=model["name"], model_url=vae["url"], token=token)
controlnets = config.get("controlnets")
if controlnets is not None:
for controlnet in controlnets:
download_controlnet(name=controlnet["name"], repo_id=controlnet["repo_id"], token=token)
loras = config.get("loras")
if loras is not None:
for lora in loras:
download_file(
url=lora["url"],
file_name=lora["name"],
file_path=BASE_CACHE_PATH_LORA,
)
textual_inversions = config.get("textual_inversions")
if textual_inversions is not None:
for textual_inversion in textual_inversions:
download_file(
url=textual_inversion["url"],
file_name=textual_inversion["name"],
file_path=BASE_CACHE_PATH_TEXTUAL_INVERSION,
)
stub = Stub("stable-diffusion-cli")
base_stub = Image.from_dockerfile(
path="Dockerfile",
context_mount=Mount.from_local_file("requirements.txt"),
)
stub.image = base_stub.extend(
dockerfile_commands=[
"FROM base",
"COPY config.yml /",
],
context_mount=Mount.from_local_file("config.yml"),
).run_function(
build_image,
secrets=[Secret.from_dotenv(__file__)],
)
app/stable_diffusion_1_5.py
+369
View File
@@ -0,0 +1,369 @@
from __future__ import annotations
import io
import os
import PIL.Image
from modal import Secret, method
from setup import (
BASE_CACHE_PATH,
BASE_CACHE_PATH_CONTROLNET,
BASE_CACHE_PATH_LORA,
BASE_CACHE_PATH_TEXTUAL_INVERSION,
stub,
)
@stub.cls(
gpu="A10G",
secrets=[Secret.from_dotenv(__file__)],
)
class SD15:
"""
SD15 is a class that runs inference using Stable Diffusion 1.5.
"""
def __enter__(self):
import diffusers
import torch
import yaml
config = {}
with open("/config.yml", "r") as file:
config = yaml.safe_load(file)
self.cache_path = os.path.join(BASE_CACHE_PATH, config["model"]["name"])
if os.path.exists(self.cache_path):
print(f"The directory '{self.cache_path}' exists.")
else:
print(f"The directory '{self.cache_path}' does not exist.")
self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
self.cache_path,
custom_pipeline="lpw_stable_diffusion",
torch_dtype=torch.float16,
use_safetensors=True,
)
# TODO: Add support for other schedulers.
self.pipe.scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
# self.pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
self.cache_path,
subfolder="scheduler",
)
# self.pipe.scheduler = diffusers.LCMScheduler.from_config(self.pipe.scheduler.config)
vae = config.get("vae")
if vae is not None:
self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
self.cache_path,
subfolder="vae",
use_safetensors=True,
)
loras = config.get("loras")
if loras is not None:
for lora in loras:
path = os.path.join(BASE_CACHE_PATH_LORA, lora["name"])
if os.path.exists(path):
print(f"The directory '{path}' exists.")
else:
print(f"The directory '{path}' does not exist. Need to execute 'modal deploy' first.")
self.pipe.load_lora_weights(".", weight_name=path)
textual_inversions = config.get("textual_inversions")
if textual_inversions is not None:
for textual_inversion in textual_inversions:
path = os.path.join(BASE_CACHE_PATH_TEXTUAL_INVERSION, textual_inversion["name"])
if os.path.exists(path):
print(f"The directory '{path}' exists.")
else:
print(f"The directory '{path}' does not exist. Need to execute 'modal deploy' first.")
self.pipe.load_textual_inversion(path)
# TODO: Repair the controlnet loading.
controlnets = config.get("controlnets")
if controlnets is not None:
for controlnet in controlnets:
path = os.path.join(BASE_CACHE_PATH_CONTROLNET, controlnet["name"])
controlnet = diffusers.ControlNetModel.from_pretrained(path, torch_dtype=torch.float16)
self.controlnet_pipe = diffusers.StableDiffusionControlNetPipeline.from_pretrained(
self.cache_path,
controlnet=controlnet,
custom_pipeline="lpw_stable_diffusion",
scheduler=self.pipe.scheduler,
vae=self.pipe.vae,
torch_dtype=torch.float16,
use_safetensors=True,
)
def _count_token(self, p: str, n: str) -> int:
"""
Count the number of tokens in the prompt and negative prompt.
"""
from transformers import CLIPTokenizer
tokenizer = CLIPTokenizer.from_pretrained(
self.cache_path,
subfolder="tokenizer",
)
token_size_p = len(tokenizer.tokenize(p))
token_size_n = len(tokenizer.tokenize(n))
token_size = token_size_p
if token_size_p <= token_size_n:
token_size = token_size_n
max_embeddings_multiples = 1
max_length = tokenizer.model_max_length - 2
if token_size > max_length:
max_embeddings_multiples = token_size // max_length + 1
print(f"token_size: {token_size}, max_embeddings_multiples: {max_embeddings_multiples}")
return max_embeddings_multiples
@method()
def run_txt2img_inference(
self,
prompt: str,
n_prompt: str,
height: int = 512,
width: int = 512,
batch_size: int = 1,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
fix_by_controlnet_tile: bool = False,
output_format: str = "png",
) -> list[bytes]:
"""
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
"""
import pillow_avif # noqa: F401
import torch
max_embeddings_multiples = self._count_token(p=prompt, n=n_prompt)
generator = torch.Generator("cuda").manual_seed(seed)
self.pipe.to("cuda")
self.pipe.enable_vae_tiling()
self.pipe.enable_xformers_memory_efficient_attention()
with torch.autocast("cuda"):
generated_images = self.pipe(
prompt=prompt * batch_size,
negative_prompt=n_prompt * batch_size,
height=height,
width=width,
num_inference_steps=steps,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
).images
base_images = generated_images
"""
Fix the generated images by the control_v11f1e_sd15_tile when `fix_by_controlnet_tile` is `True`.
https://huggingface.co/lllyasviel/control_v11f1e_sd15_tile
"""
if fix_by_controlnet_tile:
self.controlnet_pipe.to("cuda")
self.controlnet_pipe.enable_vae_tiling()
self.controlnet_pipe.enable_xformers_memory_efficient_attention()
for image in base_images:
image = self._resize_image(image=image, scale_factor=2)
with torch.autocast("cuda"):
fixed_by_controlnet = self.controlnet_pipe(
prompt=prompt * batch_size,
negative_prompt=n_prompt * batch_size,
num_inference_steps=steps,
strength=0.3,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
image=image,
).images
generated_images.extend(fixed_by_controlnet)
base_images = fixed_by_controlnet
if upscaler != "":
upscaled = self._upscale(
base_images=base_images,
half_precision=False,
tile=700,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer,
)
generated_images.extend(upscaled)
image_output = []
for image in generated_images:
with io.BytesIO() as buf:
image.save(buf, format=output_format)
image_output.append(buf.getvalue())
return image_output
@method()
def run_img2img_inference(
self,
prompt: str,
n_prompt: str,
batch_size: int = 1,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
fix_by_controlnet_tile: bool = False,
output_format: str = "png",
base_image_url: str = "",
) -> list[bytes]:
"""
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
"""
import pillow_avif # noqa: F401
import torch
from diffusers.utils import load_image
max_embeddings_multiples = self._count_token(p=prompt, n=n_prompt)
generator = torch.Generator("cuda").manual_seed(seed)
self.pipe.to("cuda")
self.pipe.enable_vae_tiling()
self.pipe.enable_xformers_memory_efficient_attention()
with torch.autocast("cuda"):
generated_images = self.pipe(
prompt=prompt * batch_size,
negative_prompt=n_prompt * batch_size,
num_inference_steps=steps,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
image=load_image(base_image_url),
).images
base_images = generated_images
"""
Fix the generated images by the control_v11f1e_sd15_tile when `fix_by_controlnet_tile` is `True`.
https://huggingface.co/lllyasviel/control_v11f1e_sd15_tile
"""
if fix_by_controlnet_tile:
self.controlnet_pipe.to("cuda")
self.controlnet_pipe.enable_vae_tiling()
self.controlnet_pipe.enable_xformers_memory_efficient_attention()
for image in base_images:
image = self._resize_image(image=image, scale_factor=2)
with torch.autocast("cuda"):
fixed_by_controlnet = self.controlnet_pipe(
prompt=prompt * batch_size,
negative_prompt=n_prompt * batch_size,
num_inference_steps=steps,
strength=0.3,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
image=image,
).images
generated_images.extend(fixed_by_controlnet)
base_images = fixed_by_controlnet
if upscaler != "":
upscaled = self._upscale(
base_images=base_images,
half_precision=False,
tile=700,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer,
)
generated_images.extend(upscaled)
image_output = []
for image in generated_images:
with io.BytesIO() as buf:
image.save(buf, format=output_format)
image_output.append(buf.getvalue())
return image_output
def _resize_image(self, image: PIL.Image.Image, scale_factor: int) -> PIL.Image.Image:
image = image.convert("RGB")
width, height = image.size
img = image.resize((width * scale_factor, height * scale_factor), resample=PIL.Image.LANCZOS)
return img
def _upscale(
self,
base_images: list[PIL.Image],
half_precision: bool = False,
tile: int = 0,
tile_pad: int = 10,
pre_pad: int = 0,
upscaler: str = "",
use_face_enhancer: bool = False,
) -> list[PIL.Image]:
"""
Upscale the generated images by the upscaler when `upscaler` is selected.
The upscaler can be selected from the following list:
- `RealESRGAN_x4plus`
- `RealESRNet_x4plus`
- `RealESRGAN_x4plus_anime_6B`
- `RealESRGAN_x2plus`
https://github.com/xinntao/Real-ESRGAN
"""
import numpy
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan import GFPGANer
from realesrgan import RealESRGANer
model_name = upscaler
if model_name == "RealESRGAN_x4plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRNet_x4plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRGAN_x4plus_anime_6B":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRGAN_x2plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
netscale = 2
else:
raise NotImplementedError("Model name not supported")
upsampler = RealESRGANer(
scale=netscale,
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", f"{model_name}.pth"),
dni_weight=None,
model=upscale_model,
tile=tile,
tile_pad=tile_pad,
pre_pad=pre_pad,
half=half_precision,
gpu_id=None,
)
if use_face_enhancer:
face_enhancer = GFPGANer(
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
upscale=netscale,
arch="clean",
channel_multiplier=2,
bg_upsampler=upsampler,
)
upscaled_imgs = []
for img in base_images:
img = numpy.array(img)
if use_face_enhancer:
_, _, enhance_result = face_enhancer.enhance(
img,
has_aligned=False,
only_center_face=False,
paste_back=True,
)
else:
enhance_result, _ = upsampler.enhance(img)
upscaled_imgs.append(PIL.Image.fromarray(enhance_result))
return upscaled_imgs
app/stable_diffusion_xl.py
+180
View File
@@ -0,0 +1,180 @@
from __future__ import annotations
import io
import os
import PIL.Image
from modal import Secret, method
from setup import BASE_CACHE_PATH, stub
@stub.cls(
gpu="A10G",
secrets=[Secret.from_dotenv(__file__)],
)
class SDXLTxt2Img:
"""
A class that wraps the Stable Diffusion pipeline and scheduler.
"""
def __enter__(self):
import diffusers
import torch
import yaml
config = {}
with open("/config.yml", "r") as file:
config = yaml.safe_load(file)
self.cache_path = os.path.join(BASE_CACHE_PATH, config["model"]["name"])
if os.path.exists(self.cache_path):
print(f"The directory '{self.cache_path}' exists.")
else:
print(f"The directory '{self.cache_path}' does not exist.")
self.pipe = diffusers.AutoPipelineForText2Image.from_pretrained(
self.cache_path,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
self.refiner_cache_path = self.cache_path + "-refiner"
self.refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_pretrained(
self.refiner_cache_path,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
@method()
def run_inference(
self,
prompt: str,
height: int = 1024,
width: int = 1024,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
output_format: str = "png",
) -> list[bytes]:
"""
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
"""
import pillow_avif # noqa
import torch
generator = torch.Generator("cuda").manual_seed(seed)
self.pipe.to("cuda")
generated_images = self.pipe(
prompt=prompt,
height=height,
width=width,
generator=generator,
).images
base_images = generated_images
for image in base_images:
self.refiner.to("cuda")
refined_images = self.refiner(
prompt=prompt,
image=image,
).images
generated_images.extend(refined_images)
base_images = refined_images
if upscaler != "":
upscaled = self._upscale(
base_images=base_images,
half_precision=False,
tile=700,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer,
)
generated_images.extend(upscaled)
image_output = []
for image in generated_images:
with io.BytesIO() as buf:
image.save(buf, format=output_format)
image_output.append(buf.getvalue())
return image_output
def _upscale(
self,
base_images: list[PIL.Image],
half_precision: bool = False,
tile: int = 0,
tile_pad: int = 10,
pre_pad: int = 0,
upscaler: str = "",
use_face_enhancer: bool = False,
) -> list[PIL.Image]:
"""
Upscale the generated images by the upscaler when `upscaler` is selected.
The upscaler can be selected from the following list:
- `RealESRGAN_x4plus`
- `RealESRNet_x4plus`
- `RealESRGAN_x4plus_anime_6B`
- `RealESRGAN_x2plus`
https://github.com/xinntao/Real-ESRGAN
"""
import numpy
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan import GFPGANer
from realesrgan import RealESRGANer
from tqdm import tqdm
model_name = upscaler
if model_name == "RealESRGAN_x4plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRNet_x4plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRGAN_x4plus_anime_6B":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
netscale = 4
elif model_name == "RealESRGAN_x2plus":
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
netscale = 2
else:
raise NotImplementedError("Model name not supported")
upsampler = RealESRGANer(
scale=netscale,
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", f"{model_name}.pth"),
dni_weight=None,
model=upscale_model,
tile=tile,
tile_pad=tile_pad,
pre_pad=pre_pad,
half=half_precision,
gpu_id=None,
)
if use_face_enhancer:
face_enhancer = GFPGANer(
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
upscale=netscale,
arch="clean",
channel_multiplier=2,
bg_upsampler=upsampler,
)
upscaled_imgs = []
for img in base_images:
img = numpy.array(img)
if use_face_enhancer:
_, _, enhance_result = face_enhancer.enhance(
img,
has_aligned=False,
only_center_face=False,
paste_back=True,
)
else:
enhance_result, _ = upsampler.enhance(img)
upscaled_imgs.append(PIL.Image.fromarray(enhance_result))
return upscaled_imgs