File size: 2,067 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 |
#!/usr/bin/env python3
#!/usr/bin/env python3
from diffusers import DiffusionPipeline
import torch
import time
import os
from pathlib import Path
from huggingface_hub import HfApi
api = HfApi()
start_time = time.time()
model_prefix = "diffusers"
pipe = DiffusionPipeline.from_pretrained(f"{model_prefix}/IF-I-IF-v1.0", torch_dtype=torch.float16, safety_checker=None, variant="fp16", use_safetensors=True)
pipe.enable_model_cpu_offload()
super_res_1_pipe = DiffusionPipeline.from_pretrained(f"{model_prefix}/IF-II-L-v1.0", text_encoder=None, safety_checker=None, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
super_res_1_pipe.enable_model_cpu_offload()
super_res_2_pipe = DiffusionPipeline.from_pretrained(f"{model_prefix}/IF-III-L-v1.0", text_encoder=None, safety_checker=None, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
super_res_2_pipe.enable_model_cpu_offload()
prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'
generator = torch.Generator("cuda").manual_seed(0)
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt").images
# save_image
pil_image = pipe.numpy_to_pil(pipe.decode_latents(image))[0]
pil_image.save(os.path.join(Path.home(), "images", "if_stage_I_0.png"))
image = super_res_1_pipe(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt", noise_level=250, num_inference_steps=50).images
# save_image
pil_image = pipe.numpy_to_pil(pipe.decode_latents(image))[0]
pil_image.save(os.path.join(Path.home(), "images", "if_stage_II_0.png"))
image = super_res_2_pipe(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, noise_level=0, num_inference_steps=40, generator=generator).images[0]
# save_image
image.save(os.path.join(Path.home(), "images", "if_stage_III_0.png"))
|