File size: 1,999 Bytes
b585c7f |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
import filelock
from diffusers import DiffusionPipeline
import torch
from src.utils import makedirs
from src.vision.sdxl import get_device
def get_pipe_make_image(gpu_id, refine=True):
device = get_device(gpu_id)
base = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
use_safetensors=True,
add_watermarker=False,
variant="fp16"
).to(device)
if not refine:
refiner = None
else:
refiner = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=base.text_encoder_2,
vae=base.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
).to(device)
return base, refiner
def make_image(prompt, filename=None, gpu_id='auto', pipe=None, guidance_scale=3.0):
if pipe is None:
base, refiner = get_pipe_make_image(gpu_id=gpu_id)
else:
base, refiner = pipe
lock_type = 'image'
base_path = os.path.join('locks', 'image_locks')
base_path = makedirs(base_path, exist_ok=True, tmp_ok=True, use_base=True)
lock_file = os.path.join(base_path, "%s.lock" % lock_type)
makedirs(os.path.dirname(lock_file)) # ensure made
with filelock.FileLock(lock_file):
# Define how many steps and what % of steps to be run on each experts (80/20) here
n_steps = 40
high_noise_frac = 0.8
# run both experts
image = base(
prompt=prompt,
num_inference_steps=n_steps,
denoising_end=high_noise_frac,
output_type="latent",
).images
image = refiner(
prompt=prompt,
num_inference_steps=n_steps,
denoising_start=high_noise_frac,
image=image,
).images[0]
if filename:
image.save(filename)
return filename
return image
|