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