EveryDream2trainer/scripts/txt2img.py

101 lines
3.7 KiB
Python
Raw Normal View History

"""
Copyright [2022] Victor C Hall
Licensed under the GNU Affero General Public License;
You may not use this code except in compliance with the License.
You may obtain a copy of the License at
https://www.gnu.org/licenses/agpl-3.0.en.html
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import random
import argparse
from torch import no_grad
from torch.cuda.amp import autocast
import torch
from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel, DDIMScheduler, DDPMScheduler, PNDMScheduler, EulerAncestralDiscreteScheduler
#from diffusers.models import AttentionBlock
from diffusers.optimization import get_scheduler
from diffusers.utils.import_utils import is_xformers_available
from transformers import CLIPTextModel, CLIPTokenizer
def __generate_sample(pipe: StableDiffusionPipeline, prompt : str, cfg: float, resolution: int, gen, steps: int = 30):
"""
generates a single sample at a given cfg scale and saves it to disk
"""
with autocast():
image = pipe(prompt,
num_inference_steps=steps,
num_images_per_prompt=1,
guidance_scale=cfg,
generator=gen,
height=resolution,
width=resolution,
).images[0]
return image
def __create_inference_pipe(unet, text_encoder, tokenizer, scheduler, vae):
"""
creates a pipeline for SD inference
"""
pipe = StableDiffusionPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=None, # save vram
requires_safety_checker=None, # avoid nag
feature_extractor=None, # must be none of no safety checker
)
return pipe
def main(args):
text_encoder = CLIPTextModel.from_pretrained(args.diffusers_path, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(args.diffusers_path, subfolder="vae")
unet = UNet2DConditionModel.from_pretrained(args.diffusers_path, subfolder="unet")
sample_scheduler = DDIMScheduler.from_pretrained(args.diffusers_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(args.diffusers_path, subfolder="tokenizer", use_fast=False)
text_encoder.eval()
vae.eval()
unet.eval()
text_encoder.to("cuda")
vae.to("cuda")
unet.to("cuda")
pipe = __create_inference_pipe(unet, text_encoder, tokenizer, sample_scheduler, vae)
seed = args.seed if args.seed != -1 else random.randint(0, 2**30)
gen = torch.Generator(device="cuda").manual_seed(seed)
img = __generate_sample(pipe, args.prompt, args.cfg_scale, args.resolution, gen=gen, steps=args.steps)
img.save(f"img_{args.prompt[0:100].replace(' ', '_')}_cfg_{args.cfg_scale}.png")
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--diffusers_path', type=str, default=None, required=True, help='path to diffusers model (from logs)')
parser.add_argument('--prompt', type=str, required=True, help='prompt to use')
parser.add_argument('--resolution', type=int, default=512, help='resolution (def: 512)')
parser.add_argument('--seed', type=int, default=-1, help='seed, use -1 for random (def: -1)')
parser.add_argument('--steps', type=int, default=50, help='inference, denoising steps (def: 50)')
parser.add_argument('--cfg_scale', type=int, default=7.5, help='unconditional guidance scale (def: 7.5)')
args = parser.parse_args()
main(args)