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

Merge pull request #100 from hodanov/feature/refactoring

Browse Source 
Replace realesrgan to the Stable Diffusion latent upscaler.
This commit is contained in:
hodanov 2024-04-21 21:36:03 +09:00 committed by GitHub
commit 6c464dc5bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 71 additions and 244 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
Makefile
View File

@ -3,12 +3,6 @@
app:
cd ./app && modal deploy __main__.py
# `--upscaler` is a name of upscaler you want to use.
# You can use upscalers the below:
# - `RealESRGAN_x4plus`
# - `RealESRNet_x4plus`
# - `RealESRGAN_x4plus_anime_6B`
# - `RealESRGAN_x2plus`
img_by_sd15_txt2img:
cd ./cmd && modal run sd15_txt2img.py \
--prompt "a photograph of an astronaut riding a horse" \
@ -17,8 +11,7 @@ img_by_sd15_txt2img:
--width 768 \
--samples 1 \
--steps 30 \
--upscaler "RealESRGAN_x2plus" \
--use-face-enhancer "False" \
--use-upscaler "True" \
--fix-by-controlnet-tile "True" \
--output-format "avif"
@ -28,8 +21,7 @@ img_by_sd15_img2img:
--n-prompt "" \
--samples 1 \
--steps 30 \
--upscaler "RealESRGAN_x2plus" \
--use-face-enhancer "False" \
--use-upscaler "True" \
--fix-by-controlnet-tile "True" \
--output-format "avif" \
--base-image-url "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"
@ -40,5 +32,4 @@ img_by_sdxl_txt2img:
--height 1024 \
--width 1024 \
--samples 1 \
--upscaler "RealESRGAN_x2plus" \
--output-format "avif"

3
README.md
View File

@ -132,8 +132,7 @@ run:
--samples 1 \
--steps 30 \
--seed 12321 |
--upscaler "RealESRGAN_x2plus" \
--use-face-enhancer "False" \
--use-upscaler "True" \
--fix-by-controlnet-tile "True" \
--output-fomart "avif"
```

5
README_ja.md
View File

@ -134,8 +134,7 @@ run:
--samples 1 \
--steps 30 \
--seed 12321 |
--upscaler "RealESRGAN_x2plus" \
--use-face-enhancer "False" \
--use-upscaler "True" \
--fix-by-controlnet-tile "True" \
--output-fomart "png"
```
@ -147,7 +146,7 @@ run:
- samples: 生成する画像の数を指定します。
- steps: ステップ数を指定します。
- seed: seedを指定します。
- upscaler: 画像の解像度を上げるためのアップスケーラーを指定します。
- use-upscaler: 画像の解像度を上げるためのアップスケーラーを有効にします。
- fix-by-controlnet-tile: ControlNet 1.1 Tileの利用有無を指定します。有効にすると、崩れた画像を修復しつつ、高解像度な画像を生成します。
- output-format: 出力フォーマットを指定します。avifも指定可能です。

8
app/Dockerfile
View File

@ -5,10 +5,4 @@ RUN apt-get update \
&& 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/cu121 --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
&& pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121 --no-cache-dir

26
app/requirements.txt
View File

@ -1,27 +1,15 @@
invisible_watermark
accelerate
diffusers[torch]==0.27.2
onnxruntime==1.17.3
safetensors==0.4.3
accelerate
torch==2.2.2
transformers==4.40.0
xformers==0.0.25.post1
realesrgan==0.3.0
basicsr>=1.4.2
facexlib>=0.3.0
gfpgan>=1.3.8
scipy==1.13.0
opencv-python
Pillow
pillow-avif-plugin
torchvision==0.17.2
tqdm
invisible_watermark # To help viewers identify the images as machine-generated.
onnxruntime==1.17.3 # ONNX Runtime uses the following optimizations to speed up Stable Diffusion in CUDA.
safetensors==0.4.3 # To store tensors safely.
controlnet_aux
pyyaml
# Use the below in 'download_from_original_stable_diffusion_ckpt'.
omegaconf==2.3.0
Pillow
pillow-avif-plugin # To save images in AVIF format.
pyyaml # To read the configuration file by written YAML.
peft

12
app/setup.py
View File

@ -9,6 +9,7 @@ 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"
BASE_CACHE_PATH_UPSCALER = "/vol/cache/upscaler"
def download_file(url, file_name, file_path):
@ -25,6 +26,15 @@ def download_file(url, file_name, file_path):
f.write(downloaded)
def download_upscaler():
"""
Download the stabilityai/sd-x2-latent-upscaler.
"""
model_id = "stabilityai/sd-x2-latent-upscaler"
upscaler = diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained(model_id)
upscaler.save_pretrained(BASE_CACHE_PATH_UPSCALER, safe_serialization=True)
def download_controlnet(name: str, repo_id: str, token: str):
"""
Download a controlnet.
@ -130,6 +140,8 @@ def build_image():
file_path=BASE_CACHE_PATH_TEXTUAL_INVERSION,
)
download_upscaler()
app = App("stable-diffusion-cli")
base_stub = Image.from_dockerfile(

131
app/stable_diffusion_1_5.py
View File

@ -10,6 +10,7 @@ from setup import (
BASE_CACHE_PATH_CONTROLNET,
BASE_CACHE_PATH_LORA,
BASE_CACHE_PATH_TEXTUAL_INVERSION,
BASE_CACHE_PATH_UPSCALER,
app,
)
@ -53,6 +54,11 @@ class SD15:
)
# self.pipe.scheduler = diffusers.LCMScheduler.from_config(self.pipe.scheduler.config)
self.upscaler = diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained(
BASE_CACHE_PATH_UPSCALER,
torch_dtype=torch.float16,
)
vae = config.get("vae")
if vae is not None:
self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
@ -133,8 +139,7 @@ class SD15:
batch_size: int = 1,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
use_upscaler: bool = False,
fix_by_controlnet_tile: bool = False,
output_format: str = "png",
) -> list[bytes]:
@ -187,16 +192,18 @@ class SD15:
generated_images.extend(fixed_by_controlnet)
base_images = fixed_by_controlnet
# TODO: Upscaler stopped working due to update of dependent packages. Replace with diffusers upscaler.
# 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)
if use_upscaler:
self.upscaler.to("cuda")
self.upscaler.enable_xformers_memory_efficient_attention()
upscaled = self.upscaler(
prompt=prompt,
negative_prompt=n_prompt,
image=base_images[0],
num_inference_steps=steps,
guidance_scale=0,
generator=generator,
).images
generated_images.extend(upscaled)
image_output = []
for image in generated_images:
@ -214,8 +221,7 @@ class SD15:
batch_size: int = 1,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
use_upscaler: bool = False,
fix_by_controlnet_tile: bool = False,
output_format: str = "png",
base_image_url: str = "",
@ -269,14 +275,17 @@ class SD15:
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,
)
if use_upscaler:
self.upscaler.to("cuda")
self.upscaler.enable_xformers_memory_efficient_attention()
upscaled = self.upscaler(
prompt=prompt,
negative_prompt=n_prompt,
image=base_images[0],
num_inference_steps=steps,
guidance_scale=0,
generator=generator,
).images
generated_images.extend(upscaled)
image_output = []
@ -292,81 +301,3 @@ class SD15:
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

99
app/stable_diffusion_xl.py
View File

@ -96,8 +96,7 @@ class SDXLTxt2Img:
width: int = 1024,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
use_upscaler: bool = False,
output_format: str = "png",
) -> list[bytes]:
"""
@ -157,15 +156,14 @@ class SDXLTxt2Img:
# 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)
# if use_upscaler:
# upscaled = self._upscale(
# base_images=base_images,
# half_precision=False,
# tile=700,
# upscaler=upscaler,
# )
# generated_images.extend(upscaled)
image_output = []
for image in generated_images:
@ -180,82 +178,3 @@ class SDXLTxt2Img:
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
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

6
cmd/sd15_img2img.py
View File

@ -15,8 +15,7 @@ def main(
batch_size: int = 1,
steps: int = 20,
seed: int = -1,
upscaler: str = "",
use_face_enhancer: str = "False",
use_upscaler: str = "False",
fix_by_controlnet_tile: str = "False",
output_format: str = "png",
base_image_url: str = "",
@ -38,8 +37,7 @@ def main(
batch_size=batch_size,
steps=steps,
seed=seed_generated,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer == "True",
use_upscaler=use_upscaler == "True",
fix_by_controlnet_tile=fix_by_controlnet_tile == "True",
output_format=output_format,
base_image_url=base_image_url,

6
cmd/sd15_txt2img.py
View File

@ -17,8 +17,7 @@ def main(
batch_size: int = 1,
steps: int = 20,
seed: int = -1,
upscaler: str = "",
use_face_enhancer: str = "False",
use_upscaler: str = "",
fix_by_controlnet_tile: str = "False",
output_format: str = "png",
):
@ -41,8 +40,7 @@ def main(
batch_size=batch_size,
steps=steps,
seed=seed_generated,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer == "True",
use_upscaler=use_upscaler == "True",
fix_by_controlnet_tile=fix_by_controlnet_tile == "True",
output_format=output_format,
)

6
cmd/sdxl_txt2img.py
View File

@ -16,8 +16,7 @@ def main(
samples: int = 5,
steps: int = 20,
seed: int = -1,
upscaler: str = "",
use_face_enhancer: str = "False",
use_upscaler: str = "False",
output_format: str = "png",
):
"""
@ -38,8 +37,7 @@ def main(
width=width,
steps=steps,
seed=seed_generated,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer == "True",
use_upscaler=use_upscaler == "True",
output_format=output_format,
)
util.save_images(directory, images, seed_generated, i, output_format)