import spaces import gradio as gr import time import torch import numpy as np from PIL import Image from segment_utils import( segment_image, restore_result, ) from diffusers import ( DiffusionPipeline, T2IAdapter, MultiAdapter, AutoencoderKL, EulerAncestralDiscreteScheduler, ) from controlnet_aux import ( CannyDetector, LineartDetector, ) BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DEFAULT_EDIT_PROMPT = "RAW photo, Fujifilm XT3, sharp hair, high resolution hair, hair tones, natural hair, magazine hair, white color hair" DEFAULT_CATEGORY = "hair" canny_detector = CannyDetector() lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators") lineart_detector = lineart_detector.to(DEVICE) adapters = MultiAdapter( [ T2IAdapter.from_pretrained( "TencentARC/t2i-adapter-lineart-sdxl-1.0", torch_dtype=torch.float16, varient="fp16", ), T2IAdapter.from_pretrained( "TencentARC/t2i-adapter-canny-sdxl-1.0", torch_dtype=torch.float16, varient="fp16", ), ] ) adapters = adapters.to(torch.float16) basepipeline = DiffusionPipeline.from_pretrained( BASE_MODEL, torch_dtype=torch.float16, variant="fp16", use_safetensors=True, vae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16), scheduler=EulerAncestralDiscreteScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler"), adapter=adapters, custom_pipeline="./pipelines/pipeline_sdxl_adapter_img2img.py", ) basepipeline = basepipeline.to(DEVICE) basepipeline.enable_model_cpu_offload() @spaces.GPU(duration=30) def image_to_image( input_image: Image, edit_prompt: str, seed: int, num_steps: int, guidance_scale: float, strength: float, generate_size: int, cond_scale1: float = 1.2, cond_scale2: float = 1.2, lineart_detect:float = 0.375, canny_detect:float = 0.375, ): run_task_time = 0 time_cost_str = '' run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str) lineart_image = lineart_detector(input_image, int(generate_size * lineart_detect), generate_size) canny_image = canny_detector(input_image, int(generate_size * canny_detect), generate_size) cond_image = [lineart_image, canny_image] cond_scale = [cond_scale1, cond_scale2] generator = torch.Generator(device=DEVICE).manual_seed(seed) generated_image = basepipeline( generator=generator, prompt=edit_prompt, image=input_image, height=generate_size, width=generate_size, guidance_scale=guidance_scale, strength=strength, num_inference_steps=num_steps, adapter_image=cond_image, adapter_conditioning_scale=cond_scale, ).images[0] run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str) return generated_image, time_cost_str def make_inpaint_condition(image, image_mask): image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0 assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size" image[image_mask > 0.5] = -1.0 # set as masked pixel image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) image = torch.from_numpy(image) return image def get_time_cost(run_task_time, time_cost_str): now_time = int(time.time()*1000) if run_task_time == 0: time_cost_str = 'start' else: if time_cost_str != '': time_cost_str += f'-->' time_cost_str += f'{now_time - run_task_time}' run_task_time = now_time return run_task_time, time_cost_str def create_demo() -> gr.Blocks: with gr.Blocks() as demo: croper = gr.State() with gr.Row(): with gr.Column(): edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT) generate_size = gr.Number(label="Generate Size", value=512) with gr.Column(): num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps") guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale") with gr.Column(): strength = gr.Slider(minimum=0, maximum=3, value=0.2, step=0.1, label="Strength") with gr.Accordion("Advanced Options", open=False): mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True) mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation") seed = gr.Number(label="Seed", value=8) category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False) cond_scale1 = gr.Slider(minimum=0, maximum=3, value=0.8, step=0.1, label="Cond_scale1") cond_scale2 = gr.Slider(minimum=0, maximum=3, value=0.3, step=0.1, label="Cond_scale2") lineart_detect = gr.Slider(minimum=0, maximum=1, value=0.375, step=0.01, label="Lineart Detect") canny_detect = gr.Slider(minimum=0, maximum=1, value=0.75, step=0.01, label="Canny Detect") g_btn = gr.Button("Edit Image") with gr.Row(): with gr.Column(): input_image = gr.Image(label="Input Image", type="pil") with gr.Column(): restored_image = gr.Image(label="Restored Image", type="pil", interactive=False) with gr.Column(): origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False) generated_image = gr.Image(label="Generated Image", type="pil", interactive=False) generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False) g_btn.click( fn=segment_image, inputs=[input_image, category, generate_size, mask_expansion, mask_dilation], outputs=[origin_area_image, croper], ).success( fn=image_to_image, inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, strength, generate_size, cond_scale1, cond_scale2, lineart_detect, canny_detect], outputs=[generated_image, generated_cost], ).success( fn=restore_result, inputs=[croper, category, generated_image], outputs=[restored_image], ) return demo