sdxl_lpw.py

from diffusers import DiffusionPipeline
from transformers import CLIPTokenizer
import torch
import os

tokenizer = CLIPTokenizer.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="tokenizer", device_map="auto"
)
# load both base & refiner
base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True,
    # FIXME: Might be a bug? Output arrays of default tokenizers have different size. 
    tokenizer=tokenizer, tokenizer_2=tokenizer
)
base.unet = torch.compile(base.unet, mode="reduce-overhead", fullgraph=True)
base.to("cuda")

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)
refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
refiner.to("cuda")

# Define how many steps and what % of steps to be run on each experts (80/20) here
n_steps = 40
high_noise_frac = 0.8

# compel = Compel(
#   tokenizer=[base.tokenizer, base.tokenizer_2] ,
#   text_encoder=[base.text_encoder, base.text_encoder_2],
#   returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
#   requires_pooled=[False, True],
#   truncate_long_prompts=False
# )
os.makedirs("output-img/sdxl_lpw", exist_ok=True)
with open("canva.benchmark.txt", "r") as f:
    for index, prompt in enumerate(f):
        image = base(
            # prompt_embeds=conditioning,
            # pooled_prompt_embeds=pooled,
            # negative_prompt_embeds=negative_conditioning,
            # negative_pooled_prompt_embeds=negative_pooled,
            prompt=prompt,
            prompt_2=prompt,
            negative_prompt="",
            negative_prompt_2="",
            num_inference_steps=n_steps,
            denoising_end=high_noise_frac,
            output_type="latent",
        ).images
        image = refiner(
            # prompt_embeds=conditioning,
            # pooled_prompt_embeds=pooled,
            # negative_prompt_embeds=negative_conditioning,
            # negative_pooled_prompt_embeds=negative_pooled,
            prompt=prompt,
            num_inference_steps=n_steps,
            denoising_start=high_noise_frac,
            image=image,
        ).images[0]
        image.save(f"output-img/sdxl/{index+1}.png")