patchedstabledifftoonnx/run-batch.py

# Copyright 2022 Dirk Moerenhout. All rights reserved.
#
# This program is free software: you can redistribute it and/or modify it under the terms
# of the GNU General Public License as published by the Free Software Foundation,
# either version 3 of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
# without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with this program. If not,
# see <https://www.gnu.org/licenses/>.

# We need sys for argv
import sys
# We need os.path for isdir, isfile
import os.path
# Our settings are in json format
import json
# To be safe we force gc to lower RAM pressure
import gc
# We want to replace the text encoder in the pipeline
import functools
# We want to parse arguments
import argparse
# Numpy is used to provide a random generator
import numpy
# We need to load images for img2img
# We want to save data to PNG
from PIL import Image, PngImagePlugin

# The pipelines
from diffusers import OnnxStableDiffusionPipeline, OnnxStableDiffusionImg2ImgPipeline
from pipeline_onnx_stable_diffusion_controlnet import OnnxStableDiffusionControlNetPipeline
# Model needed to load Text Encoder on CPU
from diffusers import OnnxRuntimeModel
# The schedulers
from diffusers import (
    DDIMScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    HeunDiscreteScheduler,
    KDPM2DiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UniPCMultistepScheduler
)

# Support special text encoders
import OnnxDiffusersUI.lpw_pipe

# Default settings
defSettings = {
    "width": 512,
    "height": 512,
    "reslist": [],
    "steps": 30,
    "stepslist": [],
    "scale": 7.5,
    "scalelist":[],
    "seed":0,
    "seedend":0,
    "seedlist":[],
    "task": "txt2img",
    "model":"sd2_1-fp16",
    "prompt": "",
    "promptlist":[],
    "negative_prompt": "",
    "textenc": "standard",
    "scheduler": "pndm",
    "schedulerlist": [],
    "strength": 0.9,
    "strengthlist": []
}

parser = argparse.ArgumentParser()

parser.add_argument(
    "--cpu-textenc",
    action="store_true",
    help="Load Text Encoder on CPU to save VRAM"
)

parser.add_argument(
    "--subdirs",
    action="store_true",
    help="Add subdirs with settings.json to projects to run"
)

parser.add_argument(
    'project',
    nargs='+',
    type=str,
    help="Provide projects as directories that contain settings.json"
)

args = parser.parse_args()

projects=args.project
if args.subdirs:
    for proj in args.project:
        obj = os.scandir(proj)
        for entry in obj:
            if entry.is_dir():
                if os.path.isfile(f"{proj}/{entry.name}/settings.json"):
                    projects.append(f"{proj}/{entry.name}")

for proj in projects:
    print("Running project "+proj)
    # Check for directory
    if os.path.isdir(proj):
        if os.path.isfile(proj+"/settings.json"):
            with open(proj+"/settings.json", encoding="utf-8") as confFile:
                projSettings=json.load(confFile)
            # Merge dictionaries with project settings taking precedence
            runSettings = defSettings | projSettings
            # We need prompts
            prereqmet=len(runSettings['prompt'])>0 or len(runSettings['promptlist'])>0
            # We need a model
            model="model/"+runSettings['model']
            prereqmet=prereqmet and os.path.isfile(model+"/unet/model.onnx")
            # We need a start image to do img2img or controlnet
            if runSettings['task']=="img2img" or runSettings['task']=="controlnet":
                infile=proj+"/input.png"
                prereqmet = prereqmet and os.path.isfile(infile)
            if prereqmet:
                sched = {
                    "ddim": DDIMScheduler.from_pretrained(model, subfolder="scheduler"),
                    "deis": DEISMultistepScheduler.from_pretrained(model, subfolder="scheduler"),
                    "dpms_ms": DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler"),
                    "dpms_ss": DPMSolverSinglestepScheduler.from_pretrained(model, subfolder="scheduler"),
                    "euler_anc": EulerAncestralDiscreteScheduler.from_pretrained(model, subfolder="scheduler"),
                    "euler": EulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler"),
                    "heun": HeunDiscreteScheduler.from_pretrained(model, subfolder="scheduler"),
                    "kdpm2": KDPM2DiscreteScheduler.from_pretrained(model, subfolder="scheduler"),
                    "lms": LMSDiscreteScheduler.from_pretrained(model, subfolder="scheduler"),
                    "pndm": PNDMScheduler.from_pretrained(model, subfolder="scheduler"),
                    "unipc": UniPCMultistepScheduler.from_pretrained(model, subfolder="scheduler")
                }
                if runSettings['task']=="img2img":
                    init_image = Image.open(infile).convert("RGB")
                    if args.cpu_textenc:
                        cputextenc=OnnxRuntimeModel.from_pretrained(model+"/text_encoder")
                        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            text_encoder=cputextenc,
                            safety_checker=None,
                            feature_extractor=None
                        )
                    else:
                        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            safety_checker=None,
                            feature_extractor=None
                        )
                elif runSettings['task']=="controlnet": 
                    init_image = Image.open(infile).convert("RGB")
                    if args.cpu_textenc:
                        cputextenc=OnnxRuntimeModel.from_pretrained(model+"/text_encoder")
                        pipe = OnnxStableDiffusionControlNetPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            text_encoder=cputextenc,
                            safety_checker=None,
                            feature_extractor=None
                        )
                    else:
                        pipe = OnnxStableDiffusionControlNetPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            safety_checker=None,
                            feature_extractor=None
                        )                 
                else:
                    if args.cpu_textenc:
                        cputextenc=OnnxRuntimeModel.from_pretrained(model+"/text_encoder")
                        pipe = OnnxStableDiffusionPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            text_encoder=cputextenc,
                            safety_checker=None,
                            feature_extractor=None
                        )
                    else:
                        pipe = OnnxStableDiffusionPipeline.from_pretrained(
                            model,
                            provider="DmlExecutionProvider",
                            revision="onnx",
                            scheduler=sched['pndm'],
                            safety_checker=None,
                            feature_extractor=None
                        )                       
                if runSettings['textenc'] == "lpw":
                    pipe._encode_prompt = functools.partial(lpw_pipe._encode_prompt, pipe)
                generator = numpy.random
                # Set schedulers for projects
                if len(runSettings['schedulerlist'])==0:
                    schedulerlist=[runSettings['scheduler']]
                else:
                    schedulerlist=runSettings['schedulerlist']
                # Set seeds for project
                if len(runSettings['seedlist'])==0:
                    if runSettings['seed']>runSettings['seedend']:
                        runSettings['seedend']=runSettings['seed']
                    seedlist=range(runSettings['seed'],runSettings['seedend']+1)
                else:
                    seedlist=runSettings['seedlist']
                # Set resolustions for project
                if len(runSettings['reslist'])==0:
                    restuples=[(runSettings['width'],runSettings['height'])]
                else:
                    restuples=[]
                    for resstr in runSettings['reslist']:
                        restuples.append(tuple(map(int, resstr.split("x"))))
                # Set steps for project
                if len(runSettings['stepslist'])==0:
                    stepslist=[runSettings['steps']]
                else:
                    stepslist=runSettings['stepslist']
                # Set guidance scales for project
                if len(runSettings['scalelist'])==0:
                    scalelist=[runSettings['scale']]
                else:
                    scalelist=runSettings['scalelist']
                # Set prompts for project
                if len(runSettings['promptlist'])==0:
                    promptlist=[runSettings['prompt']]
                else:
                    promptlist=runSettings['promptlist']
                # Set strengths for project
                if len(runSettings['strengthlist'])==0:
                    strengthlist=[runSettings['strength']]
                else:
                    strengthlist=runSettings['strengthlist']
                imgnr=len(schedulerlist)*len(promptlist)*len(seedlist)*len(restuples)*len(stepslist)*len(scalelist)*len(strengthlist)
                imgdone=0
                for scheduler in schedulerlist:
                    if not sched[scheduler]:
                        scheduler="pndm"
                    pipe.scheduler=sched[scheduler]
                    promptnum=0
                    for prompt in promptlist:
                        for seed in seedlist:
                            for res in restuples:
                                for steps in stepslist:
                                    for scale in scalelist:
                                        for strength in strengthlist:
                                            if runSettings['task']=="img2img":
                                                filename=(
                                                    f"{proj}/result-p{promptnum}-seed{seed}-{res[0]}x{res[1]}-"+
                                                    f"-steps-{steps}-{scheduler}-scale-"+str(scale).replace(".","_")+
                                                    "-strength-"+str(strength).replace(".","_")+".png"
                                                )
                                            elif runSettings['task']=="controlnet":
                                                filename=(
                                                    f"{proj}/result-p{promptnum}-seed{seed}-{res[0]}x{res[1]}-"+
                                                    f"-steps-{steps}-{scheduler}-scale-"+str(scale).replace(".","_")+
                                                    "-strength-"+str(strength).replace(".","_")+".png"      
                                                )
                                            else:
                                                filename=(
                                                    f"{proj}/result-p{promptnum}-seed{seed}-{res[0]}x{res[1]}-"+
                                                    f"-steps-{steps}-{scheduler}-scale-"+str(scale).replace(".","_")+".png"
                                                )
                                            if not os.path.isfile(filename):
                                                generator.seed(seed)
                                                if runSettings['task']=="img2img":
                                                    image = pipe(
                                                        image=init_image,
                                                        strength=strength,
                                                        prompt=prompt,
                                                        negative_prompt=runSettings['negative_prompt'],
                                                        num_inference_steps=steps,
                                                        guidance_scale=scale,
                                                        generator=generator).images[0]
                                                elif runSettings['task']=="controlnet":
                                                    image = pipe(
                                                        image=init_image,
                                                        controlnet_conditioning_scale=strength,
                                                        prompt=prompt,
                                                        negative_prompt=runSettings['negative_prompt'],
                                                        num_inference_steps=steps,
                                                        guidance_scale=scale,
                                                        generator=generator).images[0]
                                                else:
                                                    image = pipe(
                                                        prompt=prompt,
                                                        negative_prompt=runSettings['negative_prompt'],
                                                        width=res[0],
                                                        height=res[1],
                                                        num_inference_steps=steps,
                                                        guidance_scale=scale,
                                                        generator = generator).images[0]
                                                metadata = PngImagePlugin.PngInfo()
                                                metadata.add_text("Generator","Stable Diffusion ONNX https://github.com/Amblyopius/Stable-Diffusion-ONNX-FP16")
                                                metadata.add_text("SD Model (local name)",model)
                                                metadata.add_text("SD Prompt",prompt)
                                                metadata.add_text("SD Negative Prompt",runSettings['negative_prompt'])
                                                metadata.add_text("SD Scheduler",scheduler)
                                                metadata.add_text("SD Steps",str(steps))
                                                metadata.add_text("SD Guidance Scale",str(scale))
                                                image.save(filename, pnginfo = metadata)
                                            else:
                                                print("Skipping existing image!")
                                            imgdone+=1
                                            print(f"Finished {imgdone}/{imgnr}")
                        promptnum+=1
                del pipe
                gc.collect()
            else:
                print("Minimum requirements not met! Skipping")
        else:
            print("Settings not found! Skipping")
    else:
        print("Path not found! Skipping")