from diffusers import DiffusionPipeline, AutoencoderTiny from compel import Compel import torch try: import intel_extension_for_pytorch as ipex # type: ignore except: pass import psutil from config import Args from pydantic import BaseModel from PIL import Image from typing import Callable base_model = "SimianLuo/LCM_Dreamshaper_v7" WIDTH = 512 HEIGHT = 512 model_id = "wavymulder/Analog-Diffusion" lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" class Pipeline: class InputParams(BaseModel): seed: int = 2159232 prompt: str guidance_scale: float = 8.0 strength: float = 0.5 steps: int = 4 lcm_steps: int = 50 width: int = WIDTH height: int = HEIGHT @staticmethod def create_pipeline( args: Args, device: torch.device, torch_dtype: torch.dtype ) -> Callable[["Pipeline.InputParams"], Image.Image]: if args.safety_checker: pipe = DiffusionPipeline.from_pretrained(base_model) else: pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None) if args.use_taesd: pipe.vae = AutoencoderTiny.from_pretrained( "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True ) pipe.set_progress_bar_config(disable=True) pipe.to(device=device, dtype=torch_dtype) pipe.unet.to(memory_format=torch.channels_last) # Load LCM LoRA pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") # check if computer has less than 64GB of RAM using sys or os if psutil.virtual_memory().total < 64 * 1024**3: pipe.enable_attention_slicing() if args.torch_compile: pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True) pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0) compel_proc = Compel( tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, truncate_long_prompts=False, ) def predict(params: "Pipeline.InputParams") -> Image.Image: generator = torch.manual_seed(params.seed) prompt_embeds = compel_proc(params.prompt) # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps. results = pipe( prompt_embeds=prompt_embeds, generator=generator, num_inference_steps=params.steps, guidance_scale=params.guidance_scale, width=params.width, height=params.height, original_inference_steps=params.lcm_steps, output_type="pil", ) nsfw_content_detected = ( results.nsfw_content_detected[0] if "nsfw_content_detected" in results else False ) if nsfw_content_detected: return None return results.images[0] return predict