from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Optional import uuid from lcm.lcm_scheduler import LCMScheduler from lcm.lcm_pipeline import LatentConsistencyModelPipeline from lcm.lcm_i2i_pipeline import LatentConsistencyModelImg2ImgPipeline, LCMSchedulerWithTimestamp from diffusers.image_processor import PipelineImageInput # import modules.scripts as scripts # import modules.shared # from modules import script_callbacks import os import random import time import numpy as np import gradio as gr from PIL import Image, PngImagePlugin import torch scheduler = LCMScheduler.from_pretrained( "charliebaby2023/cybrpny", subfolder="scheduler") pipe = LatentConsistencyModelPipeline.from_pretrained( "charliebaby2023/cybrpny", scheduler = scheduler, safety_checker = None) # "SimianLuo/LCM_Dreamshaper_v7", scheduler = scheduler, safety_checker = None) DESCRIPTION = '''# Latent Consistency Model Running [LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7) | [Project Page](https://latent-consistency-models.github.io) | [Extension Page](https://github.com/0xbitches/sd-webui-lcm) ''' MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "768")) def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: if randomize_seed: seed = random.randint(0, MAX_SEED) return seed def save_image(img, metadata: dict): save_dir = './outputs/LCM-txt2img/' Path(save_dir).mkdir(exist_ok=True, parents=True) seed = metadata["seed"] unique_id = uuid.uuid4() filename = save_dir + f"{unique_id}-{seed}" + ".png" meta_tuples = [(k, str(v)) for k, v in metadata.items()] png_info = PngImagePlugin.PngInfo() for k, v in meta_tuples: png_info.add_text(k, v) img.save(filename, pnginfo=png_info) return filename def save_images(image_array, metadata: dict): paths = [] with ThreadPoolExecutor() as executor: paths = list(executor.map(save_image, image_array, [metadata]*len(image_array))) return paths def generate( prompt: str, seed: int = 0, width: int = 512, height: int = 512, guidance_scale: float = 8.0, num_inference_steps: int = 4, num_images: int = 4, randomize_seed: bool = False, use_fp16: bool = True, use_torch_compile: bool = False, use_cpu: bool = False, progress=gr.Progress(track_tqdm=True) ) -> Image.Image: seed = randomize_seed_fn(seed, randomize_seed) torch.manual_seed(seed) selected_device = 'cuda' if use_cpu: selected_device = "cpu" if use_fp16: use_fp16 = False print("LCM warning: running on CPU, overrode FP16 with FP32") global pipe, scheduler pipe = LatentConsistencyModelPipeline( vae= pipe.vae, text_encoder = pipe.text_encoder, tokenizer = pipe.tokenizer, unet = pipe.unet, scheduler = scheduler, safety_checker = pipe.safety_checker, feature_extractor = pipe.feature_extractor, ) # pipe = LatentConsistencyModelPipeline.from_pretrained( # "SimianLuo/LCM_Dreamshaper_v7", scheduler = scheduler, safety_checker = None) if use_fp16: pipe.to(torch_device=selected_device, torch_dtype=torch.float16) else: pipe.to(torch_device=selected_device, torch_dtype=torch.float32) # Windows does not support torch.compile for now if os.name != 'nt' and use_torch_compile: pipe.unet = torch.compile(pipe.unet, mode='max-autotune') start_time = time.time() result = pipe( prompt=prompt, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=num_images, original_inference_steps=50, output_type="pil", device = selected_device ).images paths = save_images(result, metadata={"prompt": prompt, "seed": seed, "width": width, "height": height, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps}) elapsed_time = time.time() - start_time print("LCM inference time: ", elapsed_time, "seconds") return paths, seed examples = [ "" ] with gr.Blocks() as lcm: with gr.Tab("LCM txt2img"): gr.Markdown("Latent Consistency Models: Synthesizing High-Resolution Images with Few-step Inference") gr.Markdown("Try the guide on Colab's free tier [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/R3gm/InsightSolver-Colab/blob/main/Latent_Consistency_Models.ipynb)") with gr.Row(): prompt = gr.Textbox(label="Prompt", show_label=False, lines=3, placeholder="Prompt", elem_classes=["prompt"]) run_button = gr.Button("Run", scale=0) with gr.Row(): result = gr.Gallery( label="Generated images", show_label=False, elem_id="gallery", grid=[2], preview=True ) with gr.Accordion("Advanced options", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True ) randomize_seed = gr.Checkbox( label="Randomize seed across runs", value=True) use_fp16 = gr.Checkbox( label="Run LCM in fp16 (for lower VRAM)", value=False) use_torch_compile = gr.Checkbox( label="Run LCM with torch.compile (currently not supported on Windows)", value=False) use_cpu = gr.Checkbox(label="Run LCM on CPU", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale for base", minimum=2, maximum=14, step=0.1, value=8.0, ) num_inference_steps = gr.Slider( label="Number of inference steps for base", minimum=1, maximum=8, step=1, value=4, ) with gr.Row(): num_images = gr.Slider( label="Number of images (batch count)", minimum=1, maximum=int(os.getenv("MAX_NUM_IMAGES")), step=1, value=1, ) gr.Examples( examples=examples, inputs=prompt, outputs=result, fn=generate ) run_button.click( fn=generate, inputs=[ prompt, seed, width, height, guidance_scale, num_inference_steps, num_images, randomize_seed, use_fp16, use_torch_compile, use_cpu ], outputs=[result, seed], ) if __name__ == "__main__": lcm.queue().launch()