diffusers/tools

1#!/usr/bin/env python32import torch3import os4from compel import Compel, ReturnedEmbeddingsType5from diffusers import DiffusionPipeline6from huggingface_hub import HfApi7from pathlib import Path8 9api = HfApi()10 11pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", variant="fp16", use_safetensors=True, torch_dtype=torch.float16).to("cuda")12 13compel = Compel(tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2] , text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True])14 15# upweight "ball"16prompt = ["a red cat playing with a (ball)1.5", "a red cat playing with a (ball)0.6"]17conditioning, pooled = compel(prompt)18 19 20# generate image21generator = [torch.Generator().manual_seed(33) for _ in range(len(prompt))]22images = pipeline(prompt_embeds=conditioning, pooled_prompt_embeds=pooled, generator=generator, num_inference_steps=30).images23 24for i, image in enumerate(images):25    file_name = f"bb_1_{i}"26    path = os.path.join(Path.home(), "images", f"{file_name}.png")27    image.save(path)28 29    api.upload_file(30        path_or_fileobj=path,31        path_in_repo=path.split("/")[-1],32        repo_id="patrickvonplaten/images",33        repo_type="dataset",34    )35    print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/{file_name}.png")36