diffusers/tools

1#!/usr/bin/env python32from diffusers import DiffusionPipeline, DDPMScheduler3import torch4import time5import os6from pathlib import Path7from huggingface_hub import HfApi8import random9import numpy as np10from deepfloyd_if.modules import IFStageI, IFStageII, IFStageIII, T5Embedder11import sys12 13api = HfApi()14start_time = time.time()15seed = 016use_diffusers = bool(int(sys.argv[1]))17 18t5_pos_embeds = torch.load("/home/patrick/tensors/embeds_orig.pt").to("cuda")19t5_neg_embeds = torch.load("/home/patrick/tensors/neg_embeds.pt").to("cuda")20 21def seed_everything(seed=None):22    random.seed(seed)23    os.environ['PYTHONHASHSEED'] = str(seed)24    np.random.seed(seed)25    torch.manual_seed(seed)26    torch.cuda.manual_seed(seed)27    torch.backends.cudnn.deterministic = True28    torch.backends.cudnn.benchmark = True29    return seed30 31if use_diffusers:32    pipe = DiffusionPipeline.from_pretrained("/home/patrick/if-diff-ckpts/IF-I-IF-v1.0", torch_dtype=torch.float32, use_safetensors=True, text_encoder=None, safety_checker=None)33    config = dict(pipe.scheduler.config)34    config["timestep_spacing"] = "even_border"35    pipe.scheduler = DDPMScheduler.from_config(config)36    pipe.to("cuda")37 38    with torch.no_grad():39        # text_embeddings = t5.get_text_embeddings([prompt])40        seed_everything(0)41        out_image = pipe(prompt_embeds=t5_pos_embeds, negative_prompt_embeds=t5_neg_embeds, num_inference_steps=5).images[0]42        out_image.save("/home/patrick/images/if_diff.png")43else:44    if_I = IFStageI(device="cuda", dir_or_name="/home/patrick/IF-I-IF-v1.0/", model_kwargs={"precision": "fp32"})45    if_I_kwargs = {}46    if_I_kwargs['negative_t5_embs'] = t5_neg_embeds47    if_I_kwargs['seed'] = seed48    if_I_kwargs['t5_embs'] = t5_pos_embeds49    if_I_kwargs['aspect_ratio'] = "1:1"50    if_I_kwargs['progress'] = True51    if_I_kwargs['sample_timestep_respacing'] = '5'52 53    seed_everything(0)54    stageI_generations, _ = if_I.embeddings_to_image(**if_I_kwargs)55 56    if_I.to_images(stageI_generations)[0].save("/home/patrick/images/if_ref.png")57