|
|
|
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(): |
|
|
|
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") |
|
|
