diffusers/tools

1#!/usr/bin/env python32import hf_image_uploader as hiu3import torch4from compel import Compel, ReturnedEmbeddingsType5from diffusers import DiffusionPipeline6 7pipe = DiffusionPipeline.from_pretrained(8  "stabilityai/stable-diffusion-xl-base-1.0",9  variant="fp16",10  torch_dtype=torch.float1611)12pipe2 = DiffusionPipeline.from_pretrained(13  "stabilityai/stable-diffusion-xl-refiner-1.0",14  variant="fp16",15  torch_dtype=torch.float1616)17pipe.to("cuda")18pipe2.to("cuda")19# pipe.enable_model_cpu_offload()20# pipe2.enable_model_cpu_offload()21 22compel = Compel(23  tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,24  text_encoder=[pipe.text_encoder, pipe.text_encoder_2],25  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,26  requires_pooled=[False, True]27)28 29compel2 = Compel(30  tokenizer=pipe.tokenizer_2,31  text_encoder=pipe.text_encoder_2,32  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,33  requires_pooled=True,34)35 36# apply weights37prompt = ["a red cat playing with a (ball)1.5", "a red cat playing with a (ball)0.6"]38conditioning, pooled = compel(prompt)39conditioning2, pooled2 = compel2(prompt)40 41# generate image42 43for _ in range(3):44    generator = [torch.Generator().manual_seed(i) for i in range(len(prompt))]45 46    image = pipe(prompt_embeds=conditioning, pooled_prompt_embeds=pooled, generator=generator, num_inference_steps=30, output_type="latent").images47    image = pipe2(image=image, prompt_embeds=conditioning2, pooled_prompt_embeds=pooled2, generator=generator, num_inference_steps=20).images[0]48    hiu.upload(image, "patrickvonplaten/images")49