diffusers/tools

1#!/usr/bin/env python32import torch3import numpy as np4from huggingface_hub import HfApi5 6from diffusers import ShapEPipeline7from diffusers.utils import export_to_gif8 9api = HfApi()10 11device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')12 13batch_size = 114guidance_scale = 15.015prompt = "a red table"16prompt = "A chair that looks like an avocado"17torch.manual_seed(0)18 19repo = "openai/shap-e"20pipe = ShapEPipeline.from_pretrained(repo)21pipe = pipe.to(device)22 23generator = torch.Generator(device="cuda").manual_seed(0)24 25prompts = [26    "A chair that looks like an avocado",27    "An airplane that looks like a banana",28    "A spaceship",29    "A birthday cupcake",30    "A chair that looks like a tree",31    "A green boot",32    "A penguin",33    "Ube ice cream cone",34    "A bowl of vegetables",35]36 37for prompt in prompts:38    images = pipe(39        prompt, 40        num_images_per_prompt=batch_size, 41        generator=generator, 42        guidance_scale=guidance_scale,43        num_inference_steps=64, 44        frame_size=256, 45        output_type='pil'46    ).images47 48    path = f"/home/patrick/images/{'_'.join(prompt.split())}.gif"49    export_to_gif(images[0], path)50 51 52    api.upload_file(53        path_or_fileobj=path,54        path_in_repo=path.split("/")[-1],55        repo_id="patrickvonplaten/images",56        repo_type="dataset",57    )58    print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/{path.split('/')[-1]}")59