File size: 1,692 Bytes
113dbd0 |
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 |
import PIL.Image
from diffusers import DiffusionPipeline
import torch
import numpy as np
from stable_diffusion_server.utils import log_time
pipe = DiffusionPipeline.from_pretrained(
"models/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
pipe.to("cuda")
refiner = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=pipe.text_encoder_2,
vae=pipe.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
refiner.to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
use_refiner = True
with log_time('diffuse'):
with torch.inference_mode():
image = pipe(prompt=prompt, output_type="latent" if use_refiner else "pil").images[0]
# experiment try deleting a whole bunch of pixels and see if the refiner can recreate them
# delete top 30% of pixels
# image = image[0:0.7]
#pixels to delete
# pixels_to_delete = int(0.3 * 1024)
# delete top 30% of pixels
# image.save("latent.png")
# image_data = PIL.Image.fromarray(image)
# image_data.save("latent.png")
# image = np.array(image)
pixels_to_delete = int(0.3 * image.shape[0])
idx_to_delete = np.ones(image.shape[0], dtype=bool, device="cuda")
idx_to_delete[:pixels_to_delete] = False
image[idx_to_delete] = [0,0,0]
# image_data = PIL.Image.fromarray(image)
# image_data.save("latentcleared.png")
image = refiner(prompt=prompt, image=image[None, :]).images[0]
|