Stable Diffusion (Diffusers) / Google Colab の環境でControlNet 1.1 を使ってバッチ処理で画像を作成する#3 img2img

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel,StableDiffusionControlNetImg2ImgPipeline
#

from diffusers import UniPCMultistepScheduler
from diffusers.models import AutoencoderKL
from diffusers.utils import load_image
#import torch.utils
#from controlnet_aux import PidiNetDetector,HEDdetector, MidasDetector, MLSDdetector, OpenposeDetector, PidiNetDetector, NormalBaeDetector, LineartDetector, LineartAnimeDetector, CannyDetector, ContentShuffleDetector, ZoeDetector, MediapipeFaceDetector, SamDetector, LeresDetector
from controlnet_aux.processor import Processor

from PIL import Image, ImageOps

#import cv2
#import numpy as np
from natsort import natsorted
from huggingface_hub import HfApi
from pathlib import Path

import torch
import datetime
import os
import random
import glob

#device = "cuda"
#device = "cpu"
if device == "cpu":
  torch_dtype=torch.float32
else:
  torch_dtype=torch.float16
  device = "cuda"
#
#init_img = Image.open("/content/drive/MyDrive/Images/sunflower.png")
#init_img = load_image(    "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")
file_format = "%Y%m%d_%H%M%S"
file_list1 = glob.glob(os.path.join(load_path, "*.png"))
file_list1.extend(glob.glob(os.path.join(load_path, "*.jpg")))
file_list1 = natsorted(file_list1)

file_list2 = glob.glob(os.path.join(controlnet_path, "*.png"))
file_list2.extend(glob.glob(os.path.join(controlnet_path, "*.jpg")))
file_list2 = natsorted(file_list2)
if controlnet_image_loop==False :
  file_list = file_list1
  file_listx = file_list2
else:
  file_listx = file_list1
  file_list = file_list2
if not file_list2 :
  file_list = file_list1
  file_listx = file_list2
#vae = AutoencoderKL.from_pretrained(vae)

#画像生成に使うスケジューラー
#scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
#scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")

#canny = CannyDetector()
#openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
#hed = HEDdetector.from_pretrained('lllyasviel/Annotators')
#hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
if controlnet_preprocessor_id != "":
  processor = Processor(controlnet_preprocessor_id)


#controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16)
#controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16)
#controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
#
controlnet = ControlNetModel.from_pretrained(controlnet_processor_id, torch_dtype=torch_dtype)

#パイプラインの作成
#pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id,
                                              #scheduler=scheduler,#
if controlnet_image_loop==True and not file_list1:
  pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id,
                                                      controlnet=controlnet,
                                                     # vae=vae,
                                                      #custom_pipeline="lpw_stable_diffusion",
                                                      safety_checker=None,
                                                      torch_dtype=torch_dtype)
else:
 pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(model_id,                                                      #scheduler=scheduler,
                                                      controlnet=controlnet,
                                                     # vae=vae,
                                                      #custom_pipeline="lpw_stable_diffusion",
                                                      safety_checker=None,
                                                      torch_dtype=torch_dtype)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
#
pipe.enable_model_cpu_offload()
#
pipe.to(device)

#NSFW規制を無効化する
if pipe.safety_checker is not None:
  pipe.safety_checker = lambda images, **kwargs: (images, False)


if seed is None or seed == -1:
  init_Seed = random.randint(0, 2147483647)
else:
  init_Seed = seed

def resize_image(image, target_width, target_height):
    original_width, original_height = image.size
    original_aspect_ratio = original_width / original_height
    target_aspect_ratio = target_width / target_height
    aspect_ratio_difference = original_aspect_ratio / target_aspect_ratio
    if aspect_ratio_difference > 1:
        new_width = int(target_width * aspect_ratio_difference)
        new_height = target_height
    else:
        new_width = target_width
        new_height = int(target_height / aspect_ratio_difference)
    resized_image = image.resize((new_width, new_height))
    padded_image = ImageOps.pad(resized_image, (target_width, target_height), color='white')
    return padded_image
idx2 = 0

for list_path in file_list:
  idx2 += 1
  if controlnet_image_loop==True :
    if controlnet_first_image==True :
      if idx2 > 1:
        break

  if controlnet_image_loop==False :
    img_path = list_path
    infile_name = os.path.basename(img_path)
    initfile_path =  os.path.join(controlnet_path, infile_name)
    if not os.path.exists(initfile_path) :
      initfile_path =  img_path
    if controlnet_first_image==True :
      if file_listx:
        initfile_path = file_listx[0]
  else:
    if file_listx:
      img_path = file_listx[0]
    else:
      img_path = ""
    initfile_path =  list_path
  if file_list1:
    open_img = Image.open(img_path)
    open_img = resize_image(open_img,width,height)
  #  open_img = Image.open(img_path).convert("RGB")
  #  canny_image = canny(open_img)
  #  openpose_image = openpose(open_img)
  #  scribble_image =  hed(open_img, scribble=True)
  init_img = Image.open(initfile_path)
  if controlnet_image_loop==True and not file_list1:
    pass
  else:
    init_img = resize_image(init_img,width,height)
  if controlnet_preprocessor_id != "":
    init_img =  processor(init_img)
  controlnet_save_path = f"/content/output/controlnet"

  controlnet_image_name = os.path.basename(initfile_path)
  controlnet_image_name_no_extension = os.path.splitext(controlnet_image_name)[0]
  controlnet_image_name_extension = os.path.splitext(controlnet_image_name)[1]
  controlnet_image_name = controlnet_image_name_no_extension + f".png"
  controlnet_save_pathname = os.path.join(controlnet_save_path, controlnet_image_name)
  #
  init_img.save(controlnet_save_pathname)

  image_name = os.path.basename(img_path)
  image_name_no_extension = os.path.splitext(image_name)[0]
  image_name_extension = os.path.splitext(image_name)[1]
  if controlnet_image_loop==True and not file_list1:
  # 現在の日本時間を取得
    jst_dattetime = datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=9)))
    image_name_no_extension = jst_dattetime.strftime(file_format)

  idx = 0
  while idx  < int(batch_count):
   #generator
   mSeed = init_Seed + idx
   generator = torch.Generator(device=device).manual_seed(mSeed)
  #images = []
   if controlnet_image_loop==True and not file_list1:
     image = pipe(prompt,
                image=init_img,
                negative_prompt=negative_prompt,
                width=width, height=height, generator=generator,
                controlnet_conditioning_scale=controlnet_conditioning_scale,
                guidance_scale=CFG_scale, num_inference_steps=Steps,
                guess_mode=guess_mode
                ).images[0]
   else:
     image = pipe(prompt,
                image=open_img,
                control_image=init_img,
                negative_prompt=negative_prompt,
                width=width, height=height, generator=generator,
                strength=strength,
                controlnet_conditioning_scale=controlnet_conditioning_scale,
                guidance_scale=CFG_scale, num_inference_steps=Steps,
                guess_mode=guess_mode
                ).images[0]
   #出力する画像の名前を生成する
   #outfile_name = (jst_dattetime.strftime(file_format)+ "_" + str(mSeed)+ "-" + str(idx))
   outfile_name = (image_name_no_extension+ "_" + controlnet_image_name_no_extension + "_" +  str(mSeed)+ "-" + str(idx))
   image_name = outfile_name + f".png"

   #画像を保存する
   save_pathname = os.path.join(save_path, image_name)
   image.save(save_pathname)
   idx += 1