|
|
|
|
|
import os |
|
from dotenv import load_dotenv |
|
from Helpers import name_formatter, weights_dir, capture_message |
|
from contextlib import ExitStack |
|
import torch |
|
from torch import autocast |
|
from diffusers import StableDiffusionPipeline, DDIMScheduler |
|
import base64 |
|
from io import BytesIO |
|
from PIL import Image |
|
import uuid |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
if device.type != 'cuda': |
|
raise ValueError("need to run on GPU") |
|
|
|
|
|
class Pipeline: |
|
def __init__(self, data): |
|
load_dotenv() |
|
self.output_images = os.getenv('OUTPUT_IMAGES') |
|
self.data = data |
|
|
|
def generate(self): |
|
for preset in self.data['presets']: |
|
self.run(preset) |
|
|
|
def run(self, preset): |
|
torch.cuda.empty_cache() |
|
capture_message('Pipeline: Run') |
|
prompt = preset['prompt'] % (name_formatter(self.data['name'])) |
|
n_samples = preset.get('n_samples', 1) |
|
guidance_scale = preset.get('guidance_scale', 7.5) |
|
ddim_steps = preset.get('ddim_steps', 50) |
|
ddim_eta = preset.get('ddim_eta', 0.0) |
|
n_iter = preset.get('n_iter', 1) |
|
height = preset.get('height', 512) |
|
width = preset.get('width', 512) |
|
channels = preset.get('channels', 4) |
|
scale = preset.get('scale', 7.5) |
|
seed = preset.get('seed', 7.5) |
|
preset_id = preset.get('preset_id', 1) |
|
|
|
scheduler = DDIMScheduler(beta_start=0.00085, |
|
beta_end=0.012, |
|
beta_schedule="scaled_linear", |
|
clip_sample=False, |
|
set_alpha_to_one=False) |
|
|
|
txt2img_pipe = StableDiffusionPipeline.from_pretrained( |
|
weights_dir(), |
|
scheduler=scheduler, |
|
safety_checker=None, |
|
torch_dtype=torch.float16, |
|
use_auth_token="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" |
|
) |
|
|
|
txt2img_pipe = txt2img_pipe.to(device) |
|
g_cuda = torch.Generator(device='cuda') |
|
g_cuda.manual_seed(seed) |
|
|
|
with ExitStack() as stack: |
|
if device == "cpu": |
|
_ = stack.enter_context(autocast(device)) |
|
|
|
images = txt2img_pipe( |
|
prompt=[prompt] * n_samples, |
|
guidance_scale=guidance_scale, |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
generator=g_cuda |
|
).images |
|
|
|
for img in images: |
|
img.save(self.output_images + str(preset_id) + '----' + '%s.png' % uuid.uuid4()) |
|
|
