File size: 2,054 Bytes
e58dd86 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
#!/usr/bin/env python3
from diffusers import DiffusionPipeline, DDPMScheduler
import torch
import time
import os
from pathlib import Path
from huggingface_hub import HfApi
import random
import numpy as np
from deepfloyd_if.modules import IFStageI, IFStageII, IFStageIII, T5Embedder
import sys
api = HfApi()
start_time = time.time()
seed = 0
use_diffusers = bool(int(sys.argv[1]))
t5_pos_embeds = torch.load("/home/patrick/tensors/embeds_orig.pt").to("cuda")
t5_neg_embeds = torch.load("/home/patrick/tensors/neg_embeds.pt").to("cuda")
def seed_everything(seed=None):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
return seed
if use_diffusers:
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)
config = dict(pipe.scheduler.config)
config["timestep_spacing"] = "even_border"
pipe.scheduler = DDPMScheduler.from_config(config)
pipe.to("cuda")
with torch.no_grad():
# text_embeddings = t5.get_text_embeddings([prompt])
seed_everything(0)
out_image = pipe(prompt_embeds=t5_pos_embeds, negative_prompt_embeds=t5_neg_embeds, num_inference_steps=5).images[0]
out_image.save("/home/patrick/images/if_diff.png")
else:
if_I = IFStageI(device="cuda", dir_or_name="/home/patrick/IF-I-IF-v1.0/", model_kwargs={"precision": "fp32"})
if_I_kwargs = {}
if_I_kwargs['negative_t5_embs'] = t5_neg_embeds
if_I_kwargs['seed'] = seed
if_I_kwargs['t5_embs'] = t5_pos_embeds
if_I_kwargs['aspect_ratio'] = "1:1"
if_I_kwargs['progress'] = True
if_I_kwargs['sample_timestep_respacing'] = '5'
seed_everything(0)
stageI_generations, _ = if_I.embeddings_to_image(**if_I_kwargs)
if_I.to_images(stageI_generations)[0].save("/home/patrick/images/if_ref.png")
|