|
import torch |
|
from PIL.Image import Image |
|
from diffusers import StableDiffusionXLPipeline, LCMScheduler |
|
from pipelines.models import TextToImageRequest |
|
from torch import Generator |
|
|
|
|
|
def load_pipeline() -> StableDiffusionXLPipeline: |
|
pipeline: StableDiffusionXLPipeline = StableDiffusionXLPipeline.from_pretrained( |
|
"./models/newdream-sdxl-20", |
|
torch_dtype=torch.float16, |
|
local_files_only=True, |
|
).to("cuda") |
|
adapter_id = "./models/lora" |
|
pipeline.scheduler = LCMScheduler.from_config({ |
|
"_class_name": "LCMScheduler", |
|
"_diffusers_version": "0.29.0", |
|
"beta_end": 0.012, |
|
"beta_schedule": "scaled_linear", |
|
"beta_start": 0.00085, |
|
"clip_sample": False, |
|
"clip_sample_range": 1.0, |
|
"num_train_timesteps": 1000, |
|
"original_inference_steps": 300, |
|
"prediction_type": "epsilon", |
|
"rescale_betas_zero_snr": False, |
|
"sample_max_value": 1.0, |
|
"set_alpha_to_one": False, |
|
"steps_offset": 1, |
|
"thresholding": False, |
|
"timestep_scaling": 10.0, |
|
"timestep_spacing": "leading", |
|
"trained_betas": None, |
|
}) |
|
|
|
|
|
pipeline.load_lora_weights(adapter_id) |
|
|
|
pipeline(prompt="") |
|
|
|
return pipeline |
|
|
|
|
|
def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image: |
|
generator = Generator(pipeline.device).manual_seed(request.seed) if request.seed else None |
|
|
|
return pipeline( |
|
prompt=request.prompt, |
|
negative_prompt=request.negative_prompt, |
|
width=request.width, |
|
height=request.height, |
|
generator=generator, |
|
num_inference_steps=4, |
|
guidance_scale=1.5, |
|
).images[0] |
|
|
