+ {/if}
++ There are {currentQueueSize} + user(s) sharing the same GPU, affecting real-time performance. Maximum queue size is {maxQueueSize}. + Duplicate and run it on your own GPU. +
+ {/if} +Loading...
++ This demo showcases + LCM LoRA + ControlNet + Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet" + title: str = "LCM + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 2, min=1, max=6, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=2, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.8, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ).to(device) + + self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + base_model, + safety_checker=None, + controlnet=controlnet_canny, + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + print("\nRunning sfast compile\n") + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if args.onediff: + print("\nRunning onediff compile\n") + from onediff.infer_compiler import oneflow_compile + + self.pipe.unet = oneflow_compile(self.pipe.unet) + self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder) + self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder) + self.pipe.controlnet = oneflow_compile(self.pipe.controlnet) + + self.canny_torch = SobelOperator(device=device) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(params.prompt) + prompt = None + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + control_image=control_image, + prompt_embeds=prompt_embeds, + prompt=prompt, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetDepthFlashSD.py b/server/pipelines/controlnetDepthFlashSD.py new file mode 100644 index 0000000000000000000000000000000000000000..9a63705533411c7bd4ba99c2bec515e813317708 --- /dev/null +++ b/server/pipelines/controlnetDepthFlashSD.py @@ -0,0 +1,264 @@ +from diffusers import ( + StableDiffusionControlNetImg2ImgPipeline, + ControlNetModel, + AutoencoderTiny, + LCMScheduler, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from transformers import pipeline + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +controlnet_model = "lllyasviel/control_v11f1p_sd15_depth" +model_id = "runwayml/stable-diffusion-v1-5" +taesd_model = "madebyollin/taesd" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + LCM LoRA ++ ControlNet + Image to Imasge pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+loras+sd15" + title: str = "LCM + LoRA + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + base_model_id: str = Field( + "plasmo/woolitize", + title="Base Model", + values=list(base_models.keys()), + field="select", + id="base_model_id", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 768, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 768, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=2, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.8, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ).to(device) + + self.pipes = {} + + for base_model_id in base_models.keys(): + pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + base_model_id, + safety_checker=None, + controlnet=controlnet_canny, + ) + self.pipes[base_model_id] = pipe + + self.canny_torch = SobelOperator(device=device) + + for pipe in self.pipes.values(): + pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) + pipe.set_progress_bar_config(disable=True) + if device.type != "mps": + pipe.unet.to(memory_format=torch.channels_last) + + if args.taesd: + pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + # Load LCM LoRA + pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") + pipe.to(device=device, dtype=torch_dtype).to(device) + if args.compel: + self.compel_proc = Compel( + tokenizer=pipe.tokenizer, + text_encoder=pipe.text_encoder, + truncate_long_prompts=False, + ) + 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", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + pipe = self.pipes[params.base_model_id] + + activation_token = base_models[params.base_model_id] + prompt = f"{activation_token} {params.prompt}" + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(prompt) + prompt = None + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetLoraSD15QRCode.py b/server/pipelines/controlnetLoraSD15QRCode.py new file mode 100644 index 0000000000000000000000000000000000000000..0c4f56840c04ffe966c4b191439894bf31a6eb69 --- /dev/null +++ b/server/pipelines/controlnetLoraSD15QRCode.py @@ -0,0 +1,226 @@ +from diffusers import ( + StableDiffusionControlNetImg2ImgPipeline, + ControlNetModel, + LCMScheduler, + 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, Field +from PIL import Image +import math + +taesd_model = "madebyollin/taesd" +controlnet_model = "monster-labs/control_v1p_sd15_qrcode_monster" +base_model = "nitrosocke/mo-di-diffusion" +lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" +default_prompt = "abstract art of a men with curly hair by Pablo Picasso" +page_content = """ ++ This demo showcases + LCM LoRA ++ ControlNet + Image to Imasge pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+loras+sd15" + title: str = "LCM + LoRA + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 5, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=2, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.6, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + blend: float = Field( + 0.1, + min=0.0, + max=1.0, + step=0.001, + title="Blend", + field="range", + hide=True, + id="blend", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_qrcode = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype, subfolder="v2" + ).to(device) + + if args.safety_checker: + self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + base_model, + controlnet=controlnet_qrcode, + ) + else: + self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + base_model, + safety_checker=None, + controlnet=controlnet_qrcode, + ) + + self.control_image = Image.open("qr-code.png").convert("RGB").resize((512, 512)) + + self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.pipe.set_progress_bar_config(disable=True) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + # Load LCM LoRA + self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") + self.pipe.to(device=device, dtype=torch_dtype).to(device) + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (512, 512))], + control_image=[Image.new("RGB", (512, 512))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + + prompt = f"modern disney style {params.prompt}" + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(prompt) + prompt = None + + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + blend_qr_image = Image.blend( + params.image, self.control_image, alpha=params.blend + ) + results = self.pipe( + image=blend_qr_image, + control_image=self.control_image, + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + + return results.images[0] diff --git a/server/pipelines/controlnetLoraSDXL-Lightning.py b/server/pipelines/controlnetLoraSDXL-Lightning.py new file mode 100644 index 0000000000000000000000000000000000000000..c19cbdd221dd2365582ba5069c8d4b97cd0ecbcb --- /dev/null +++ b/server/pipelines/controlnetLoraSDXL-Lightning.py @@ -0,0 +1,296 @@ +from diffusers import ( + UNet2DConditionModel, + StableDiffusionXLControlNetImg2ImgPipeline, + ControlNetModel, + AutoencoderKL, + AutoencoderTiny, + EulerDiscreteScheduler, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from pipelines.utils.canny_gpu import SobelOperator +from huggingface_hub import hf_hub_download +from safetensors.torch import load_file + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +controlnet_model = "diffusers/controlnet-canny-sdxl-1.0-small" +base = "stabilityai/stable-diffusion-xl-base-1.0" +repo = "ByteDance/SDXL-Lightning" +ckpt = "sdxl_lightning_2step_unet.safetensors" +taesd_model = "madebyollin/taesdxl" +NUM_STEPS = 2 + + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + LCM LoRA ++ SDXL + Controlnet + Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+loras+sdxl+lightning" + title: str = "SDXL + LCM + LoRA + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=2.0, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 1, + min=0.25, + max=1.0, + step=0.0001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.5, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + if args.taesd: + vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ) + else: + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ) + + unet = UNet2DConditionModel.from_config(base, subfolder="unet") + unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device.type)) + self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( + base, + unet=unet, + torch_dtype=torch_dtype, + variant="fp16", + controlnet=controlnet_canny, + vae=vae, + ) + + # Ensure sampler uses "trailing" timesteps. + self.pipe.scheduler = EulerDiscreteScheduler.from_config( + self.pipe.scheduler.config, timestep_spacing="trailing" + ) + + self.canny_torch = SobelOperator(device=device) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image | None: + generator = torch.manual_seed(params.seed) + + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = NUM_STEPS + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=NUM_STEPS, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetLoraSDXL.py b/server/pipelines/controlnetLoraSDXL.py new file mode 100644 index 0000000000000000000000000000000000000000..44356542debff813412ef9ec4c6428ceca679e1f --- /dev/null +++ b/server/pipelines/controlnetLoraSDXL.py @@ -0,0 +1,294 @@ +from diffusers import ( + StableDiffusionXLControlNetImg2ImgPipeline, + ControlNetModel, + LCMScheduler, + AutoencoderKL, + AutoencoderTiny, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from pipelines.utils.canny_gpu import SobelOperator + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +controlnet_model = "diffusers/controlnet-canny-sdxl-1.0" +model_id = "stabilityai/stable-diffusion-xl-base-1.0" +lcm_lora_id = "latent-consistency/lcm-lora-sdxl" +taesd_model = "madebyollin/taesdxl" + + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + LCM LoRA ++ SDXL + Controlnet + Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+loras+sdxl" + title: str = "SDXL + LCM + LoRA + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=10, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=2.0, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 1, + min=0.25, + max=1.0, + step=0.0001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.5, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ).to(device) + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( + model_id, + safety_checker=None, + controlnet=controlnet_canny, + vae=vae, + ) + self.canny_torch = SobelOperator(device=device) + # Load LCM LoRA + self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") + self.pipe.load_lora_weights( + "CiroN2022/toy-face", + weight_name="toy_face_sdxl.safetensors", + adapter_name="toy", + ) + self.pipe.set_adapters(["lcm", "toy"], adapter_weights=[1.0, 0.8]) + + self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetMistoLineHyperSDXL.py b/server/pipelines/controlnetMistoLineHyperSDXL.py new file mode 100644 index 0000000000000000000000000000000000000000..098829f29180d0f584ffc7d6e8294fdb7237a1c3 --- /dev/null +++ b/server/pipelines/controlnetMistoLineHyperSDXL.py @@ -0,0 +1,288 @@ +from diffusers import ( + StableDiffusionXLControlNetImg2ImgPipeline, + ControlNetModel, + AutoencoderKL, + TCDScheduler, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from controlnet_aux import AnylineDetector +from huggingface_hub import hf_hub_download + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +# controlnet_model = "diffusers/controlnet-canny-sdxl-1.0" +controlnet_model = "TheMistoAI/MistoLine" +model_id = "stabilityai/stable-diffusion-xl-base-1.0" +taesd_model = "madebyollin/taesdxl" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SD Turbo +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+sd15Turbo" + title: str = "SDv1.5 Turbo + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 4402026899276587, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.21, + min=0, + max=10, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.8, + min=0.10, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.325, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ) + self.pipes = {} + + self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + base_model, + controlnet=controlnet_canny, + safety_checker=None, + torch_dtype=torch_dtype, + ) + + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + print("\nRunning sfast compile\n") + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if args.onediff: + print("\nRunning onediff compile\n") + from onediff.infer_compiler import oneflow_compile + + self.pipe.unet = oneflow_compile(self.pipe.unet) + self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder) + self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder) + self.pipe.controlnet = oneflow_compile(self.pipe.controlnet) + + self.canny_torch = SobelOperator(device=device) + + self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.compel: + from compel import Compel + + self.pipe.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=True, + ) + + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt = params.prompt + prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + results = self.pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetSDXLTurbo.py b/server/pipelines/controlnetSDXLTurbo.py new file mode 100644 index 0000000000000000000000000000000000000000..60adafcd72de8f4b6cb652c8af388281e4ebb04f --- /dev/null +++ b/server/pipelines/controlnetSDXLTurbo.py @@ -0,0 +1,281 @@ +from diffusers import ( + StableDiffusionXLControlNetImg2ImgPipeline, + ControlNetModel, + AutoencoderKL, + AutoencoderTiny, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from pipelines.utils.canny_gpu import SobelOperator + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +controlnet_model = "diffusers/controlnet-canny-sdxl-1.0" +model_id = "stabilityai/sdxl-turbo" +taesd_model = "madebyollin/taesdxl" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SDXL Turbo +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+SDXL+Turbo" + title: str = "SDXL Turbo + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 2, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=10, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.5, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, torch_dtype=torch_dtype + ).to(device) + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + + self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( + model_id, + safety_checker=None, + controlnet=controlnet_canny, + vae=vae, + ) + self.canny_torch = SobelOperator(device=device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype).to(device) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/controlnetSegmindVegaRT.py b/server/pipelines/controlnetSegmindVegaRT.py new file mode 100644 index 0000000000000000000000000000000000000000..559485d626cb7a0ed692468d80ef18c0bb2021ca --- /dev/null +++ b/server/pipelines/controlnetSegmindVegaRT.py @@ -0,0 +1,289 @@ +from diffusers import ( + StableDiffusionXLControlNetImg2ImgPipeline, + ControlNetModel, + AutoencoderKL, + AutoencoderTiny, + LCMScheduler, +) +from compel import Compel, ReturnedEmbeddingsType +import torch +from pipelines.utils.canny_gpu import SobelOperator + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +controlnet_model = "diffusers/controlnet-canny-sdxl-1.0" +base_model = "segmind/Segmind-Vega" +lora_model = "segmind/Segmind-VegaRT" +taesd_model = "madebyollin/taesdxl" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + Segmind-VegaRT +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet+SegmindVegaRT" + title: str = "SegmindVegaRT + Controlnet" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 2, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=1, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + controlnet_scale: float = Field( + 0.5, + min=0, + max=1.0, + step=0.001, + title="Controlnet Scale", + field="range", + hide=True, + id="controlnet_scale", + ) + controlnet_start: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet Start", + field="range", + hide=True, + id="controlnet_start", + ) + controlnet_end: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Controlnet End", + field="range", + hide=True, + id="controlnet_end", + ) + canny_low_threshold: float = Field( + 0.31, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 0.125, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + controlnet_canny = ControlNetModel.from_pretrained( + controlnet_model, + torch_dtype=torch_dtype, + ).to(device) + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( + base_model, + safety_checker=None, + controlnet=controlnet_canny, + vae=vae, + ) + self.canny_torch = SobelOperator(device=device) + + self.pipe.load_lora_weights(lora_model) + self.pipe.fuse_lora() + self.pipe.scheduler = LCMScheduler.from_pretrained( + base_model, subfolder="scheduler" + ) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype).to(device) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + control_image=[Image.new("RGB", (768, 768))], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + control_image = self.canny_torch( + params.image, params.canny_low_threshold, params.canny_high_threshold + ) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + control_image=control_image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + controlnet_conditioning_scale=params.controlnet_scale, + control_guidance_start=params.controlnet_start, + control_guidance_end=params.controlnet_end, + ) + + result_image = results.images[0] + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = control_image.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + + return result_image diff --git a/server/pipelines/img2img.py b/server/pipelines/img2img.py new file mode 100644 index 0000000000000000000000000000000000000000..679bd198e7f630b71b5ee5fbcd460ff1edb4c2f0 --- /dev/null +++ b/server/pipelines/img2img.py @@ -0,0 +1,171 @@ +from diffusers import ( + AutoPipelineForImage2Image, + 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, Field +from PIL import Image +import math + +base_model = "SimianLuo/LCM_Dreamshaper_v7" +taesd_model = "madebyollin/taesd" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +page_content = """ ++ This demo showcases + LCM +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image LCM" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 4, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 768, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 768, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.2, + min=0, + max=20, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base_model, + safety_checker=None, + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(params.prompt) + prompt = None + + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/img2imgSDTurbo.py b/server/pipelines/img2imgSDTurbo.py new file mode 100644 index 0000000000000000000000000000000000000000..d3f935128d8ffbe8943ad7ae9af5a7d753076fab --- /dev/null +++ b/server/pipelines/img2imgSDTurbo.py @@ -0,0 +1,187 @@ +from diffusers import ( + AutoPipelineForImage2Image, + AutoencoderTiny, +) +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + + +base_model = "stabilityai/sd-turbo" +taesd_model = "madebyollin/taesd" + +default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SDXL Turbo +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image SDXL" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base_model, + safety_checker=None, + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + print("\nRunning sfast compile\n") + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if args.onediff: + print("\nRunning onediff compile\n") + from onediff.infer_compiler import oneflow_compile + + self.pipe.unet = oneflow_compile(self.pipe.unet) + self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder) + self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + if args.compel: + from compel import Compel + + self.pipe.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=True, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + prompt = params.prompt + prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + + results = self.pipe( + image=params.image, + prompt_embeds=prompt_embeds, + prompt=prompt, + negative_prompt=params.negative_prompt, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=1.1, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/img2imgSDXL-Lightning.py b/server/pipelines/img2imgSDXL-Lightning.py new file mode 100644 index 0000000000000000000000000000000000000000..1a0a9613972366f86f2df3dbf766159dce26cf1f --- /dev/null +++ b/server/pipelines/img2imgSDXL-Lightning.py @@ -0,0 +1,216 @@ +from diffusers import ( + AutoPipelineForImage2Image, + AutoencoderTiny, + AutoencoderKL, + UNet2DConditionModel, + EulerDiscreteScheduler, +) +from compel import Compel, ReturnedEmbeddingsType +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +from safetensors.torch import load_file +from huggingface_hub import hf_hub_download +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +base = "stabilityai/stable-diffusion-xl-base-1.0" +repo = "ByteDance/SDXL-Lightning" +ckpt = "sdxl_lightning_2step_unet.safetensors" +taesd_model = "madebyollin/taesdxl" +NUM_STEPS = 2 + +default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SDXL Turbo +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image SDXL-Lightning" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=10, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=1, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + if args.taesd: + vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ) + else: + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + + unet = UNet2DConditionModel.from_config(base, subfolder="unet") + unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device.type)) + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base, + unet=unet, + torch_dtype=torch_dtype, + variant="fp16", + safety_checker=False, + vae=vae, + ) + # Ensure sampler uses "trailing" timesteps. + self.pipe.scheduler = EulerDiscreteScheduler.from_config( + self.pipe.scheduler.config, timestep_spacing="trailing" + ) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/img2imgSDXLTurbo.py b/server/pipelines/img2imgSDXLTurbo.py new file mode 100644 index 0000000000000000000000000000000000000000..13815b5913e6f2527d789e2725eac197783540d9 --- /dev/null +++ b/server/pipelines/img2imgSDXLTurbo.py @@ -0,0 +1,194 @@ +from diffusers import ( + AutoPipelineForImage2Image, + AutoencoderTiny, +) +from compel import Compel, ReturnedEmbeddingsType +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +base_model = "stabilityai/sdxl-turbo" +taesd_model = "madebyollin/taesdxl" + +default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SDXL Turbo +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image SDXL" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=10, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 768, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 768, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=1, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base_model, + safety_checker=None, + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/img2imgSDXS512.py b/server/pipelines/img2imgSDXS512.py new file mode 100644 index 0000000000000000000000000000000000000000..1b7d23a4eb6d7d8145da0cb032e50c23ff131b3d --- /dev/null +++ b/server/pipelines/img2imgSDXS512.py @@ -0,0 +1,162 @@ +from diffusers import AutoPipelineForImage2Image, 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, Field +from PIL import Image +import math + +base_model = "IDKiro/sdxs-512-0.9" +taesd_model = "madebyollin/taesd" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" +page_content = """ ++ This demo showcases + LCM +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image SDXS" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=20, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base_model, + safety_checker=None, + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(params.prompt) + prompt = None + + results = self.pipe( + image=params.image, + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + strength=params.strength, + num_inference_steps=params.steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + return results.images[0] diff --git a/server/pipelines/img2imgSegmindVegaRT.py b/server/pipelines/img2imgSegmindVegaRT.py new file mode 100644 index 0000000000000000000000000000000000000000..4523b60d643e07f316192b1de3e81576e5c631c0 --- /dev/null +++ b/server/pipelines/img2imgSegmindVegaRT.py @@ -0,0 +1,202 @@ +from diffusers import ( + AutoPipelineForImage2Image, + LCMScheduler, + AutoencoderTiny, +) +from compel import Compel, ReturnedEmbeddingsType +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +import math + +base_model = "segmind/Segmind-Vega" +lora_model = "segmind/Segmind-VegaRT" +taesd_model = "madebyollin/taesdxl" + +default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + SegmindVegaRT +Image to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image Playground 256" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 1, min=1, max=10, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.0, + min=0, + max=1, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + strength: float = Field( + 0.5, + min=0.25, + max=1.0, + step=0.001, + title="Strength", + field="range", + hide=True, + id="strength", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = AutoPipelineForImage2Image.from_pretrained( + base_model, + safety_checker=None, + variant="fp16", + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + self.pipe.load_lora_weights(lora_model) + self.pipe.fuse_lora() + self.pipe.scheduler = LCMScheduler.from_pretrained( + base_model, subfolder="scheduler" + ) + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + print("Running torch compile") + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=False + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=False + ) + + self.pipe( + prompt="warmup", + image=[Image.new("RGB", (768, 768))], + ) + if args.compel: + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + steps = params.steps + strength = params.strength + if int(steps * strength) < 1: + steps = math.ceil(1 / max(0.10, strength)) + + results = self.pipe( + image=params.image, + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + strength=strength, + num_inference_steps=steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/pix2pix/__init__.py b/server/pipelines/pix2pix/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/server/pipelines/pix2pix/model.py b/server/pipelines/pix2pix/model.py new file mode 100644 index 0000000000000000000000000000000000000000..35095850692f08217f7b297dac80ac1875f2a91e --- /dev/null +++ b/server/pipelines/pix2pix/model.py @@ -0,0 +1,59 @@ +# https://github.com/GaParmar/img2img-turbo/blob/main/src/model.py +from diffusers import DDPMScheduler + + +def make_1step_sched(): + noise_scheduler_1step = DDPMScheduler.from_pretrained( + "stabilityai/sd-turbo", subfolder="scheduler" + ) + noise_scheduler_1step.set_timesteps(1, device="cuda") + noise_scheduler_1step.alphas_cumprod = noise_scheduler_1step.alphas_cumprod.cuda() + return noise_scheduler_1step + + +def my_vae_encoder_fwd(self, sample): + sample = self.conv_in(sample) + l_blocks = [] + # down + for down_block in self.down_blocks: + l_blocks.append(sample) + sample = down_block(sample) + # middle + sample = self.mid_block(sample) + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + self.current_down_blocks = l_blocks + return sample + + +def my_vae_decoder_fwd(self, sample, latent_embeds=None): + sample = self.conv_in(sample) + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + # middle + sample = self.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + if not self.ignore_skip: + skip_convs = [ + self.skip_conv_1, + self.skip_conv_2, + self.skip_conv_3, + self.skip_conv_4, + ] + # up + for idx, up_block in enumerate(self.up_blocks): + skip_in = skip_convs[idx](self.incoming_skip_acts[::-1][idx] * self.gamma) + # add skip + sample = sample + skip_in + sample = up_block(sample, latent_embeds) + else: + for idx, up_block in enumerate(self.up_blocks): + sample = up_block(sample, latent_embeds) + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + return sample diff --git a/server/pipelines/pix2pix/pix2pix_turbo.py b/server/pipelines/pix2pix/pix2pix_turbo.py new file mode 100644 index 0000000000000000000000000000000000000000..c476c7bdce814e23ccb2f580895c8f6727ed227a --- /dev/null +++ b/server/pipelines/pix2pix/pix2pix_turbo.py @@ -0,0 +1,213 @@ +# https://github.com/GaParmar/img2img-turbo/blob/main/src/pix2pix_turbo.py +import os +import requests +import sys +import pdb +import copy +from tqdm import tqdm +import torch +from transformers import AutoTokenizer, PretrainedConfig, CLIPTextModel +from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler +from diffusers.utils.peft_utils import set_weights_and_activate_adapters +from peft import LoraConfig + +from pipelines.pix2pix.model import ( + make_1step_sched, + my_vae_encoder_fwd, + my_vae_decoder_fwd, +) + + +class TwinConv(torch.nn.Module): + def __init__(self, convin_pretrained, convin_curr): + super(TwinConv, self).__init__() + self.conv_in_pretrained = copy.deepcopy(convin_pretrained) + self.conv_in_curr = copy.deepcopy(convin_curr) + self.r = None + + def forward(self, x): + x1 = self.conv_in_pretrained(x).detach() + x2 = self.conv_in_curr(x) + return x1 * (1 - self.r) + x2 * (self.r) + + +class Pix2Pix_Turbo(torch.nn.Module): + def __init__(self, name, ckpt_folder="checkpoints"): + super().__init__() + self.tokenizer = AutoTokenizer.from_pretrained( + "stabilityai/sd-turbo", subfolder="tokenizer" + ) + self.text_encoder = CLIPTextModel.from_pretrained( + "stabilityai/sd-turbo", subfolder="text_encoder" + ).cuda() + self.sched = make_1step_sched() + + vae = AutoencoderKL.from_pretrained("stabilityai/sd-turbo", subfolder="vae") + unet = UNet2DConditionModel.from_pretrained( + "stabilityai/sd-turbo", subfolder="unet" + ) + + if name == "edge_to_image": + url = "https://www.cs.cmu.edu/~img2img-turbo/models/edge_to_image_loras.pkl" + os.makedirs(ckpt_folder, exist_ok=True) + outf = os.path.join(ckpt_folder, "edge_to_image_loras.pkl") + if not os.path.exists(outf): + print(f"Downloading checkpoint to {outf}") + response = requests.get(url, stream=True) + total_size_in_bytes = int(response.headers.get("content-length", 0)) + block_size = 1024 # 1 Kibibyte + progress_bar = tqdm( + total=total_size_in_bytes, unit="iB", unit_scale=True + ) + with open(outf, "wb") as file: + for data in response.iter_content(block_size): + progress_bar.update(len(data)) + file.write(data) + progress_bar.close() + if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes: + print("ERROR, something went wrong") + print(f"Downloaded successfully to {outf}") + p_ckpt = outf + sd = torch.load(p_ckpt, map_location="cpu") + unet_lora_config = LoraConfig( + r=sd["rank_unet"], + init_lora_weights="gaussian", + target_modules=sd["unet_lora_target_modules"], + ) + + if name == "sketch_to_image_stochastic": + # download from url + url = "https://www.cs.cmu.edu/~img2img-turbo/models/sketch_to_image_stochastic_lora.pkl" + os.makedirs(ckpt_folder, exist_ok=True) + outf = os.path.join(ckpt_folder, "sketch_to_image_stochastic_lora.pkl") + if not os.path.exists(outf): + print(f"Downloading checkpoint to {outf}") + response = requests.get(url, stream=True) + total_size_in_bytes = int(response.headers.get("content-length", 0)) + block_size = 1024 # 1 Kibibyte + progress_bar = tqdm( + total=total_size_in_bytes, unit="iB", unit_scale=True + ) + with open(outf, "wb") as file: + for data in response.iter_content(block_size): + progress_bar.update(len(data)) + file.write(data) + progress_bar.close() + if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes: + print("ERROR, something went wrong") + print(f"Downloaded successfully to {outf}") + p_ckpt = outf + sd = torch.load(p_ckpt, map_location="cpu") + unet_lora_config = LoraConfig( + r=sd["rank_unet"], + init_lora_weights="gaussian", + target_modules=sd["unet_lora_target_modules"], + ) + convin_pretrained = copy.deepcopy(unet.conv_in) + unet.conv_in = TwinConv(convin_pretrained, unet.conv_in) + + vae.encoder.forward = my_vae_encoder_fwd.__get__( + vae.encoder, vae.encoder.__class__ + ) + vae.decoder.forward = my_vae_decoder_fwd.__get__( + vae.decoder, vae.decoder.__class__ + ) + # add the skip connection convs + vae.decoder.skip_conv_1 = torch.nn.Conv2d( + 512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False + ).cuda() + vae.decoder.skip_conv_2 = torch.nn.Conv2d( + 256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False + ).cuda() + vae.decoder.skip_conv_3 = torch.nn.Conv2d( + 128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False + ).cuda() + vae.decoder.skip_conv_4 = torch.nn.Conv2d( + 128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False + ).cuda() + vae_lora_config = LoraConfig( + r=sd["rank_vae"], + init_lora_weights="gaussian", + target_modules=sd["vae_lora_target_modules"], + ) + vae.decoder.ignore_skip = False + vae.add_adapter(vae_lora_config, adapter_name="vae_skip") + unet.add_adapter(unet_lora_config) + _sd_unet = unet.state_dict() + for k in sd["state_dict_unet"]: + _sd_unet[k] = sd["state_dict_unet"][k] + unet.load_state_dict(_sd_unet) + unet.enable_xformers_memory_efficient_attention() + _sd_vae = vae.state_dict() + for k in sd["state_dict_vae"]: + _sd_vae[k] = sd["state_dict_vae"][k] + vae.load_state_dict(_sd_vae) + unet.to("cuda") + vae.to("cuda") + unet.eval() + vae.eval() + self.unet, self.vae = unet, vae + self.vae.decoder.gamma = 1 + self.timesteps = torch.tensor([999], device="cuda").long() + self.last_prompt = "" + self.caption_enc = None + self.device = "cuda" + + @torch.no_grad() + def forward(self, c_t, prompt, deterministic=True, r=1.0, noise_map=1.0): + # encode the text prompt + if prompt != self.last_prompt: + caption_tokens = self.tokenizer( + prompt, + max_length=self.tokenizer.model_max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).input_ids.cuda() + caption_enc = self.text_encoder(caption_tokens)[0] + self.caption_enc = caption_enc + self.last_prompt = prompt + + if deterministic: + encoded_control = ( + self.vae.encode(c_t).latent_dist.sample() + * self.vae.config.scaling_factor + ) + model_pred = self.unet( + encoded_control, + self.timesteps, + encoder_hidden_states=self.caption_enc, + ).sample + x_denoised = self.sched.step( + model_pred, self.timesteps, encoded_control, return_dict=True + ).prev_sample + self.vae.decoder.incoming_skip_acts = self.vae.encoder.current_down_blocks + output_image = ( + self.vae.decode(x_denoised / self.vae.config.scaling_factor).sample + ).clamp(-1, 1) + else: + # scale the lora weights based on the r value + self.unet.set_adapters(["default"], weights=[r]) + set_weights_and_activate_adapters(self.vae, ["vae_skip"], [r]) + encoded_control = ( + self.vae.encode(c_t).latent_dist.sample() + * self.vae.config.scaling_factor + ) + # combine the input and noise + unet_input = encoded_control * r + noise_map * (1 - r) + self.unet.conv_in.r = r + unet_output = self.unet( + unet_input, + self.timesteps, + encoder_hidden_states=self.caption_enc, + ).sample + self.unet.conv_in.r = None + x_denoised = self.sched.step( + unet_output, self.timesteps, unet_input, return_dict=True + ).prev_sample + self.vae.decoder.incoming_skip_acts = self.vae.encoder.current_down_blocks + self.vae.decoder.gamma = r + output_image = ( + self.vae.decode(x_denoised / self.vae.config.scaling_factor).sample + ).clamp(-1, 1) + return output_image diff --git a/server/pipelines/pix2pixTurbo.py b/server/pipelines/pix2pixTurbo.py new file mode 100644 index 0000000000000000000000000000000000000000..b2ad8853204330df8034e56109398c20f31b6065 --- /dev/null +++ b/server/pipelines/pix2pixTurbo.py @@ -0,0 +1,135 @@ +import torch +from torchvision import transforms + +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +from pipelines.pix2pix.pix2pix_turbo import Pix2Pix_Turbo +from pipelines.utils.canny_gpu import ScharrOperator + +default_prompt = "close-up photo of the joker" +page_content = """ ++ This demo showcases + One-Step Image Translation with Text-to-Image Models + +
++ Web app + Real-Time Latent Consistency Models + +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "img2img" + title: str = "Image-to-Image SDXL" + description: str = "Generates an image from a text prompt" + input_mode: str = "image" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + noise_r: float = Field( + 1.0, + min=0.01, + max=3.0, + step=0.001, + title="Noise R", + field="range", + hide=True, + id="noise_r", + ) + + deterministic: bool = Field( + True, + hide=True, + title="Deterministic", + field="checkbox", + id="deterministic", + ) + canny_low_threshold: float = Field( + 0.0, + min=0, + max=1.0, + step=0.001, + title="Canny Low Threshold", + field="range", + hide=True, + id="canny_low_threshold", + ) + canny_high_threshold: float = Field( + 1.0, + min=0, + max=1.0, + step=0.001, + title="Canny High Threshold", + field="range", + hide=True, + id="canny_high_threshold", + ) + debug_canny: bool = Field( + False, + title="Debug Canny", + field="checkbox", + hide=True, + id="debug_canny", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.model = Pix2Pix_Turbo("edge_to_image") + self.canny_torch = ScharrOperator(device=device) + self.device = device + self.last_time = 0.0 + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + canny_pil, canny_tensor = self.canny_torch( + params.image, + params.canny_low_threshold, + params.canny_high_threshold, + output_type="pil,tensor", + ) + torch.manual_seed(params.seed) + noise = torch.randn( + (1, 4, params.width // 8, params.height // 8), device=self.device + ) + canny_tensor = torch.cat((canny_tensor, canny_tensor, canny_tensor), dim=1) + output_image = self.model( + canny_tensor, + params.prompt, + params.deterministic, + params.noise_r, + noise, + ) + output_pil = transforms.ToPILImage()(output_image[0].cpu() * 0.5 + 0.5) + + result_image = output_pil + if params.debug_canny: + # paste control_image on top of result_image + w0, h0 = (200, 200) + control_image = canny_pil.resize((w0, h0)) + w1, h1 = result_image.size + result_image.paste(control_image, (w1 - w0, h1 - h0)) + return result_image diff --git a/server/pipelines/txt2img.py b/server/pipelines/txt2img.py new file mode 100644 index 0000000000000000000000000000000000000000..34f640e42bd3938d19ad48ca197b4bd6a2616d44 --- /dev/null +++ b/server/pipelines/txt2img.py @@ -0,0 +1,141 @@ +from diffusers import DiffusionPipeline, AutoencoderTiny +from compel import Compel +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +from config import Args +from pydantic import BaseModel, Field +from PIL import Image +from typing import List + +base_model = "SimianLuo/LCM_Dreamshaper_v7" +taesd_model = "madebyollin/taesd" + +default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece" + +page_content = """+ This demo showcases + LCM +Image to Image pipeline using + Diffusers with a MJPEG stream server +
++ Change the prompt to generate different images, accepts Compel syntax. +
""" + + +class Pipeline: + class Info(BaseModel): + name: str = "txt2img" + title: str = "Text-to-Image LCM" + description: str = "Generates an image from a text prompt" + input_mode: str = "text" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 4, min=2, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 768, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 768, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 8.0, + min=1, + max=30, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype) + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0) + + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(params.prompt) + prompt = None + + results = self.pipe( + prompt_embeds=prompt_embeds, + prompt=prompt, + generator=generator, + num_inference_steps=params.steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/txt2imgLora.py b/server/pipelines/txt2imgLora.py new file mode 100644 index 0000000000000000000000000000000000000000..3168b65694ceaf1bca4e2bd12004fe3ab347f646 --- /dev/null +++ b/server/pipelines/txt2imgLora.py @@ -0,0 +1,151 @@ +from diffusers import DiffusionPipeline, AutoencoderTiny, LCMScheduler +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, Field +from PIL import Image + +base_model = "wavymulder/Analog-Diffusion" +lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" +taesd_model = "madebyollin/taesd" + +default_prompt = "Analog style photograph of young Harrison Ford as Han Solo, star wars behind the scenes" + +page_content = """ ++ This demo showcases + LCM +Image to Image pipeline using + Diffusers with a MJPEG stream server. Featuring Analog-Diffusion +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "controlnet" + title: str = "Text-to-Image LCM + LoRa" + description: str = "Generates an image from a text prompt" + input_mode: str = "text" + page_content: str = page_content + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + seed: int = Field( + 8638236174640251, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 4, min=2, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 0.2, + min=0, + max=4, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + self.pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") + self.pipe.to(device=device, dtype=torch_dtype) + + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + + self.pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if args.compel: + self.compel_proc = Compel( + tokenizer=self.pipe.tokenizer, + text_encoder=self.pipe.text_encoder, + truncate_long_prompts=False, + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + prompt_embeds = None + prompt = params.prompt + if hasattr(self, "compel_proc"): + prompt_embeds = self.compel_proc(params.prompt) + prompt = None + + results = self.pipe( + prompt=prompt, + prompt_embeds=prompt_embeds, + generator=generator, + num_inference_steps=params.steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/txt2imgLoraSDXL.py b/server/pipelines/txt2imgLoraSDXL.py new file mode 100644 index 0000000000000000000000000000000000000000..d609f887c40628f4dbe86afc3d8b5b99bd6d4f7b --- /dev/null +++ b/server/pipelines/txt2imgLoraSDXL.py @@ -0,0 +1,182 @@ +from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderKL, AutoencoderTiny +from compel import Compel, ReturnedEmbeddingsType +import torch + +try: + import intel_extension_for_pytorch as ipex # type: ignore +except: + pass + +import psutil +from config import Args +from pydantic import BaseModel, Field +from PIL import Image + +model_id = "stabilityai/stable-diffusion-xl-base-1.0" +lcm_lora_id = "latent-consistency/lcm-lora-sdxl" +taesd_model = "madebyollin/taesdxl" + + +default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux" +default_negative_prompt = "blurry, low quality, render, 3D, oversaturated" +page_content = """ ++ This demo showcases + LCM LoRA + Text to Image pipeline using + Diffusers with a MJPEG stream server. +
++ Change the prompt to generate different images, accepts Compel syntax. +
+""" + + +class Pipeline: + class Info(BaseModel): + name: str = "LCM+Lora+SDXL" + title: str = "Text-to-Image SDXL + LCM + LoRA" + description: str = "Generates an image from a text prompt" + page_content: str = page_content + input_mode: str = "text" + + class InputParams(BaseModel): + prompt: str = Field( + default_prompt, + title="Prompt", + field="textarea", + id="prompt", + ) + negative_prompt: str = Field( + default_negative_prompt, + title="Negative Prompt", + field="textarea", + id="negative_prompt", + hide=True, + ) + seed: int = Field( + 2159232, min=0, title="Seed", field="seed", hide=True, id="seed" + ) + steps: int = Field( + 4, min=1, max=15, title="Steps", field="range", hide=True, id="steps" + ) + width: int = Field( + 1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width" + ) + height: int = Field( + 1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height" + ) + guidance_scale: float = Field( + 1.0, + min=0, + max=20, + step=0.001, + title="Guidance Scale", + field="range", + hide=True, + id="guidance_scale", + ) + + def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): + vae = AutoencoderKL.from_pretrained( + "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype + ) + + self.pipe = DiffusionPipeline.from_pretrained( + model_id, + safety_checker=None, + vae=vae, + ) + # Load LCM LoRA + self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm") + self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.pipe.set_progress_bar_config(disable=True) + self.pipe.to(device=device, dtype=torch_dtype).to(device) + + if args.sfast: + from sfast.compilers.stable_diffusion_pipeline_compiler import ( + compile, + CompilationConfig, + ) + + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.pipe = compile(self.pipe, config=config) + + if device.type != "mps": + self.pipe.unet.to(memory_format=torch.channels_last) + + self.pipe.compel_proc = Compel( + tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2], + text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=[False, True], + ) + if args.taesd: + self.pipe.vae = AutoencoderTiny.from_pretrained( + taesd_model, torch_dtype=torch_dtype, use_safetensors=True + ).to(device) + + if args.torch_compile: + self.pipe.unet = torch.compile( + self.pipe.unet, mode="reduce-overhead", fullgraph=True + ) + self.pipe.vae = torch.compile( + self.pipe.vae, mode="reduce-overhead", fullgraph=True + ) + self.pipe( + prompt="warmup", + ) + + def predict(self, params: "Pipeline.InputParams") -> Image.Image: + generator = torch.manual_seed(params.seed) + + prompt = params.prompt + negative_prompt = params.negative_prompt + prompt_embeds = None + pooled_prompt_embeds = None + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if hasattr(self.pipe, "compel_proc"): + _prompt_embeds, pooled_prompt_embeds = self.pipe.compel_proc( + [params.prompt, params.negative_prompt] + ) + prompt = None + negative_prompt = None + prompt_embeds = _prompt_embeds[0:1] + pooled_prompt_embeds = pooled_prompt_embeds[0:1] + negative_prompt_embeds = _prompt_embeds[1:2] + negative_pooled_prompt_embeds = pooled_prompt_embeds[1:2] + + results = self.pipe( + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + generator=generator, + num_inference_steps=params.steps, + guidance_scale=params.guidance_scale, + width=params.width, + height=params.height, + output_type="pil", + ) + + return results.images[0] diff --git a/server/pipelines/utils/canny_gpu.py b/server/pipelines/utils/canny_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..33fe74287e83db46b3a75452675f90b5d9406c28 --- /dev/null +++ b/server/pipelines/utils/canny_gpu.py @@ -0,0 +1,117 @@ +import torch +import torch.nn as nn +from torchvision.transforms import ToTensor, ToPILImage +from PIL import Image + + +class SobelOperator(nn.Module): + SOBEL_KERNEL_X = torch.tensor( + [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]] + ) + SOBEL_KERNEL_Y = torch.tensor( + [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]] + ) + + def __init__(self, device="cuda"): + super(SobelOperator, self).__init__() + self.device = device + self.edge_conv_x = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to( + self.device + ) + self.edge_conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to( + self.device + ) + self.edge_conv_x.weight = nn.Parameter( + self.SOBEL_KERNEL_X.view((1, 1, 3, 3)).to(self.device) + ) + self.edge_conv_y.weight = nn.Parameter( + self.SOBEL_KERNEL_Y.view((1, 1, 3, 3)).to(self.device) + ) + + @torch.no_grad() + def forward( + self, + image: Image.Image, + low_threshold: float, + high_threshold: float, + output_type="pil", + ) -> Image.Image | torch.Tensor | tuple[Image.Image, torch.Tensor]: + # Convert PIL image to PyTorch tensor + image_gray = image.convert("L") + image_tensor = ToTensor()(image_gray).unsqueeze(0).to(self.device) + + # Compute gradients + edge_x = self.edge_conv_x(image_tensor) + edge_y = self.edge_conv_y(image_tensor) + edge = torch.sqrt(torch.square(edge_x) + torch.square(edge_y)) + + # Apply thresholding + edge.div_(edge.max()) # Normalize to 0-1 (in-place operation) + edge[edge >= high_threshold] = 1.0 + edge[edge <= low_threshold] = 0.0 + + # Convert the result back to a PIL image + if output_type == "pil": + return ToPILImage()(edge.squeeze(0).cpu()) + elif output_type == "tensor": + return edge + elif output_type == "pil,tensor": + return ToPILImage()(edge.squeeze(0).cpu()), edge + + +class ScharrOperator(nn.Module): + SCHARR_KERNEL_X = torch.tensor( + [[-3.0, 0.0, 3.0], [-10.0, 0.0, 10.0], [-3.0, 0.0, 3.0]] + ) + SCHARR_KERNEL_Y = torch.tensor( + [[-3.0, -10.0, -3.0], [0.0, 0.0, 0.0], [3.0, 10.0, 3.0]] + ) + + def __init__(self, device="cuda"): + super(ScharrOperator, self).__init__() + self.device = device + self.edge_conv_x = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to( + self.device + ) + self.edge_conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to( + self.device + ) + self.edge_conv_x.weight = nn.Parameter( + self.SCHARR_KERNEL_X.view((1, 1, 3, 3)).to(self.device) + ) + self.edge_conv_y.weight = nn.Parameter( + self.SCHARR_KERNEL_Y.view((1, 1, 3, 3)).to(self.device) + ) + + @torch.no_grad() + def forward( + self, + image: Image.Image, + low_threshold: float, + high_threshold: float, + output_type="pil", + invert: bool = False, + ) -> Image.Image | torch.Tensor | tuple[Image.Image, torch.Tensor]: + # Convert PIL image to PyTorch tensor + image_gray = image.convert("L") + image_tensor = ToTensor()(image_gray).unsqueeze(0).to(self.device) + + # Compute gradients + edge_x = self.edge_conv_x(image_tensor) + edge_y = self.edge_conv_y(image_tensor) + edge = torch.abs(edge_x) + torch.abs(edge_y) + + # Apply thresholding + edge.div_(edge.max()) # Normalize to 0-1 (in-place operation) + edge[edge >= high_threshold] = 1.0 + edge[edge <= low_threshold] = 0.0 + if invert: + edge = 1 - edge + + # Convert the result back to a PIL image + if output_type == "pil": + return ToPILImage()(edge.squeeze(0).cpu()) + elif output_type == "tensor": + return edge + elif output_type == "pil,tensor": + return ToPILImage()(edge.squeeze(0).cpu()), edge diff --git a/server/pipelines/utils/safety_checker.py b/server/pipelines/utils/safety_checker.py new file mode 100644 index 0000000000000000000000000000000000000000..8e2cea48cb9907f588fda9e6087b00b2b5835654 --- /dev/null +++ b/server/pipelines/utils/safety_checker.py @@ -0,0 +1,169 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn +from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel +from PIL import Image + + +def cosine_distance(image_embeds, text_embeds): + normalized_image_embeds = nn.functional.normalize(image_embeds) + normalized_text_embeds = nn.functional.normalize(text_embeds) + return torch.mm(normalized_image_embeds, normalized_text_embeds.t()) + + +class StableDiffusionSafetyChecker(PreTrainedModel): + config_class = CLIPConfig + + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + self.vision_model = CLIPVisionModel(config.vision_config) + self.visual_projection = nn.Linear( + config.vision_config.hidden_size, config.projection_dim, bias=False + ) + + self.concept_embeds = nn.Parameter( + torch.ones(17, config.projection_dim), requires_grad=False + ) + self.special_care_embeds = nn.Parameter( + torch.ones(3, config.projection_dim), requires_grad=False + ) + + self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False) + self.special_care_embeds_weights = nn.Parameter( + torch.ones(3), requires_grad=False + ) + + @torch.no_grad() + def forward(self, clip_input, images): + pooled_output = self.vision_model(clip_input)[1] # pooled_output + image_embeds = self.visual_projection(pooled_output) + + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + special_cos_dist = ( + cosine_distance(image_embeds, self.special_care_embeds) + .cpu() + .float() + .numpy() + ) + cos_dist = ( + cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy() + ) + + result = [] + batch_size = image_embeds.shape[0] + for i in range(batch_size): + result_img = { + "special_scores": {}, + "special_care": [], + "concept_scores": {}, + "bad_concepts": [], + } + + # increase this value to create a stronger `nfsw` filter + # at the cost of increasing the possibility of filtering benign images + adjustment = 0.0 + + for concept_idx in range(len(special_cos_dist[0])): + concept_cos = special_cos_dist[i][concept_idx] + concept_threshold = self.special_care_embeds_weights[concept_idx].item() + result_img["special_scores"][concept_idx] = round( + concept_cos - concept_threshold + adjustment, 3 + ) + if result_img["special_scores"][concept_idx] > 0: + result_img["special_care"].append( + {concept_idx, result_img["special_scores"][concept_idx]} + ) + adjustment = 0.01 + + for concept_idx in range(len(cos_dist[0])): + concept_cos = cos_dist[i][concept_idx] + concept_threshold = self.concept_embeds_weights[concept_idx].item() + result_img["concept_scores"][concept_idx] = round( + concept_cos - concept_threshold + adjustment, 3 + ) + if result_img["concept_scores"][concept_idx] > 0: + result_img["bad_concepts"].append(concept_idx) + + result.append(result_img) + + has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result] + + return has_nsfw_concepts + + @torch.no_grad() + def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor): + pooled_output = self.vision_model(clip_input)[1] # pooled_output + image_embeds = self.visual_projection(pooled_output) + + special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds) + cos_dist = cosine_distance(image_embeds, self.concept_embeds) + + # increase this value to create a stronger `nsfw` filter + # at the cost of increasing the possibility of filtering benign images + adjustment = 0.0 + + special_scores = ( + special_cos_dist - self.special_care_embeds_weights + adjustment + ) + # special_scores = special_scores.round(decimals=3) + special_care = torch.any(special_scores > 0, dim=1) + special_adjustment = special_care * 0.01 + special_adjustment = special_adjustment.unsqueeze(1).expand( + -1, cos_dist.shape[1] + ) + + concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment + # concept_scores = concept_scores.round(decimals=3) + has_nsfw_concepts = torch.any(concept_scores > 0, dim=1) + + images[has_nsfw_concepts] = 0.0 # black image + + return images, has_nsfw_concepts + + +class SafetyChecker: + def __init__(self, device="cuda"): + from transformers import CLIPFeatureExtractor + + self.device = device + self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( + "CompVis/stable-diffusion-safety-checker" + ).to(device) + self.feature_extractor = CLIPFeatureExtractor.from_pretrained( + "openai/clip-vit-base-patch32" + ) + + def __call__( + self, images: list[Image.Image] | Image.Image + ) -> tuple[list[Image.Image], list[bool]] | tuple[Image.Image, bool]: + images_list = [images] if isinstance(images, Image.Image) else images + + safety_checker_input = self.feature_extractor( + images_list, return_tensors="pt" + ).to(self.device) + + has_nsfw_concepts = self.safety_checker( + images=[images_list], + clip_input=safety_checker_input.pixel_values.to(self.device), + ) + + if isinstance(images, Image.Image): + return images, has_nsfw_concepts[0] + + return images, has_nsfw_concepts diff --git a/server/requirements.txt b/server/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..ce927f34858edc8bf2735dccd8028c7b715e33bf --- /dev/null +++ b/server/requirements.txt @@ -0,0 +1,21 @@ +diffusers==0.28.2 +transformers==4.41.1 +--extra-index-url https://download.pytorch.org/whl/cu121; +torch==2.2.2 +fastapi==0.111.0 +uvicorn[standard]==0.30.0 +Pillow==10.3.0 +accelerate==0.30.1 +compel==2.0.2 +controlnet-aux==0.0.9 +peft==0.11.1 +xformers; sys_platform != 'darwin' or platform_machine != 'arm64' +markdown2 +safetensors +stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v1.0.5/stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64' +oneflow @ https://github.com/siliconflow/oneflow_releases/releases/download/community_cu121/oneflow-0.9.1.dev20240515+cu121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64' +onediff @ git+https://github.com/siliconflow/onediff.git@main#egg=onediff ; sys_platform != 'darwin' or platform_machine != 'arm64' +setuptools +mpmath==1.3.0 +numpy==1.* +controlnet-aux \ No newline at end of file diff --git a/server/util.py b/server/util.py new file mode 100644 index 0000000000000000000000000000000000000000..5cd80baea28668296a3578dd67b356f755382c7f --- /dev/null +++ b/server/util.py @@ -0,0 +1,40 @@ +from importlib import import_module +from types import ModuleType +from PIL import Image +import io + + +def get_pipeline_class(pipeline_name: str) -> ModuleType: + try: + module = import_module(f"pipelines.{pipeline_name}") + except ModuleNotFoundError: + raise ValueError(f"Pipeline {pipeline_name} module not found") + + pipeline_class = getattr(module, "Pipeline", None) + + if pipeline_class is None: + raise ValueError(f"'Pipeline' class not found in module '{pipeline_name}'.") + + return pipeline_class + + +def bytes_to_pil(image_bytes: bytes) -> Image.Image: + image = Image.open(io.BytesIO(image_bytes)) + return image + + +def pil_to_frame(image: Image.Image) -> bytes: + frame_data = io.BytesIO() + image.save(frame_data, format="JPEG", quality=80, optimize=True, progressive=True) + frame_data = frame_data.getvalue() + return ( + b"--frame\r\n" + + b"Content-Type: image/jpeg\r\n" + + f"Content-Length: {len(frame_data)}\r\n\r\n".encode() + + frame_data + + b"\r\n" + ) + + +def is_firefox(user_agent: str) -> bool: + return "Firefox" in user_agent