diffusers/tools

1#!/usr/bin/env python32from diffusers import DiffusionPipeline3import torch4torch.backends.cudnn.deterministic = False5torch.backends.cuda.matmul.allow_tf32 = False6torch.backends.cudnn.allow_tf32 = False7torch.backends.cudnn.benchmark = True8torch.backends.cuda.enable_flash_sdp(False)9 10# vae = AutoEncoderKL.from_pretrained("stabilityai/sdxl-vae", torch_dtype=torch.float16)11# base_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")12base_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")13base_pipe.to("cuda") # OR, pipe.enable_sequential_cpu_offload() OR,14 15# Reproducibility.16torch_seed = 420242042017refiner_seed = 69800856918refiner_strength = 0.5019prompt = "happy child flying a kite on a sunny day"20negative_prompt = ''21# Batch size.22batch_size = 223do_latent = True24 25# We're going to schedule 20 steps, and complete 50% of them using either model.26total_num_steps = 2027# We need multiple Generators.28generator = torch.Generator(device="cuda").manual_seed(torch_seed)29 30pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=base_pipe.vae, text_encoder_2=base_pipe.text_encoder_2, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")31# Using channels last layout.32pipe.unet.to(memory_format=torch.channels_last)33pipe.to("cuda") # OR, pipe.enable_sequential_cpu_offload() OR, 34 35pre_image = base_pipe(prompt=prompt, generator=generator,36         num_inference_steps=total_num_steps, negative_prompt=negative_prompt, num_images_per_prompt=batch_size, output_type="latent" if do_latent else "pil").images37 38# Generate a range from 0.1 to 0.9, with 0.1 increments.39test_strengths = [0.5]40for refiner_strength in test_strengths:41    # Generate a new set of random states for each image.42    generator_two = torch.Generator(device="cuda").manual_seed(refiner_seed)43    # Put through the refiner now.44    images = pipe(prompt=prompt, image=pre_image, aesthetic_score=10, negative_aesthetic_score=2.4, generator=generator_two,45                num_inference_steps=total_num_steps, num_images_per_prompt=batch_size, strength=refiner_strength, negative_prompt=negative_prompt).images # denoising_start46    for idx in range(0, len(images)):47        print(f'Image: {idx}')48        images[idx].save(f'/home/patrick/images/refiner_bug/test-{refiner_strength}-{idx}--{batch_size}--{do_latent}.png', format='PNG')49