diffusers/tools

1#!/usr/bin/env python32from diffusers import DiffusionPipeline, StableDiffusionPipeline, KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler, HeunDiscreteScheduler, DDIMScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, PNDMScheduler, LMSDiscreteScheduler, DPMSolverMultistepScheduler3import torch4import os5 6seed = 337inference_steps = 258 9#old_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base", custom_pipeline="/home/patrick_huggingface_co/diffusers/examples/community/sd_text2img_k_diffusion.py")10#old_pipe = old_pipe.to("cuda")11#old_pipe.set_progress_bar_config(disable=True)12pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)13#pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base", torch_dtype=torch.float16)14pipe = pipe.to("cuda")15 16 17for prompt in ["astronaut riding horse", "whale falling from sky", "magical forest", "highly photorealistic picture of johnny depp"]:18    for sampler in ["sample_dpm_2_ancestral", "euler_ancestral", "sample_dpm_2", "sample_heun", "lms", "ddim", "euler", "pndm", "dpm"]:19#    for sampler in ["sample_dpm_2_ancestral"]:20#        old_pipe.set_sampler(sampler)21#        torch.manual_seed(0)22#        image = old_pipe(prompt, height=512, width=512, num_inference_steps=inference_steps).images[0]23        folder = f"a_{'_'.join(prompt.split())}"24        os.makedirs(f"/home/patrick_huggingface_co/images/{folder}", exist_ok=True)25#        image.save(f"/home/patrick_huggingface_co/images/{folder}/{sampler}.png")26 27#        pipe = StableDiffusionPipeline(**old_pipe.components)28#        pipe = pipe.to("cuda")29#        pipe.set_progress_bar_config(disable=True)30        if sampler == "sample_dpm_2":31            pipe.scheduler = KDPM2DiscreteScheduler.from_config(pipe.scheduler.config)32        elif sampler == "sample_dpm_2_ancestral":33            pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)34        elif sampler == "sample_heun":35            pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)36        elif sampler == "ddim":37            pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)38        elif sampler == "dpm":39            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)40        elif sampler == "euler":41            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)42        elif sampler == "euler_ancestral":43            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)44        elif sampler == "pndm":45            pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)46        elif sampler == "lms":47            pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)48 49        torch.manual_seed(0)50        image = pipe(prompt, num_inference_steps=inference_steps).images[0]51 52        image.save(f"/home/patrick_huggingface_co/images/{folder}/hf_{sampler}.png")53    break54