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): | |
device = get_device(gpu_id) | |
pipe = DiffusionPipeline.from_pretrained( | |
"playgroundai/playground-v2-1024px-aesthetic", | |
torch_dtype=torch.float16, | |
use_safetensors=True, | |
add_watermarker=False, | |
variant="fp16" | |
).to(device) | |
return pipe | |
def make_image(prompt, filename=None, gpu_id='auto', pipe=None, guidance_scale=3.0): | |
if pipe is None: | |
pipe = get_pipe_make_image(gpu_id=gpu_id) | |
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): | |
image = pipe(prompt=prompt, guidance_scale=guidance_scale).images[0] | |
if filename: | |
image.save(filename) | |
return filename | |
return image | |