app.py

import gradio as gr
import spaces
import torch
from clip_slider_pipeline import CLIPSliderXL
from diffusers import StableDiffusionXLPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline, EulerDiscreteScheduler,  AutoencoderKL
import time
import numpy as np
import cv2
from PIL import Image
from ledits.pipeline_leditspp_stable_diffusion_xl import LEditsPPPipelineStableDiffusionXL

def HWC3(x):
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y

def process_controlnet_img(image):
    controlnet_img = np.array(image)
    controlnet_img = cv2.Canny(controlnet_img, 100, 200)
    controlnet_img = HWC3(controlnet_img)
    controlnet_img = Image.fromarray(controlnet_img)

# load pipelines
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", vae=vae).to("cuda", torch.float16)
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
clip_slider = CLIPSliderXL(pipe, device=torch.device("cuda"))

pipe_adapter = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash").to("cuda", torch.float16)
pipe_adapter.scheduler = EulerDiscreteScheduler.from_config(pipe_adapter.scheduler.config)
#pipe_adapter.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
# scale = 0.8
# pipe_adapter.set_ip_adapter_scale(scale)
clip_slider_ip = CLIPSliderXL(sd_pipe=pipe_adapter, device=torch.device("cuda"))

controlnet = ControlNetModel.from_pretrained(
    "xinsir/controlnet-canny-sdxl-1.0", # insert here your choice of controlnet
    torch_dtype=torch.float16
)
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe_controlnet = StableDiffusionXLControlNetPipeline.from_pretrained(
    "sd-community/sdxl-flash",
    controlnet=controlnet,
    vae=vae,
    torch_dtype=torch.float16,
)
clip_slider_controlnet = CLIPSliderXL(sd_pipe=pipe_controlnet,device=torch.device("cuda"))

pipe_inv = LEditsPPPipelineStableDiffusionXL.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", vae=vae,
    torch_dtype=torch.float16
)
clip_slider_inv = CLIPSliderXL(sd_pipe=pipe_inv,device=torch.device("cuda"))

@spaces.GPU(duration=120)
def generate(slider_x, slider_y, prompt, seed, iterations, steps, guidance_scale,
             x_concept_1, x_concept_2, y_concept_1, y_concept_2, 
             avg_diff_x_1, avg_diff_x_2,
             avg_diff_y_1, avg_diff_y_2,
             img2img_type = None, img = None, 
             controlnet_scale= None, ip_adapter_scale=None,
             edit_threshold=None, edit_guidance_scale = None,
             init_latents=None, zs=None):
    
    start_time = time.time()
    # check if avg diff for directions need to be re-calculated
    print("slider_x", slider_x)
    print("x_concept_1", x_concept_1, "x_concept_2", x_concept_2)
    
    if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]):
        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], num_iterations=iterations)
        avg_diff_0 = avg_diff[0].to(torch.float16)
        avg_diff_1 = avg_diff[1].to(torch.float16)
        x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
    
    print("avg_diff_0", avg_diff_0.dtype)
    
    if not sorted(slider_y) == sorted([y_concept_1, y_concept_2]):
        avg_diff_2nd = clip_slider.find_latent_direction(slider_y[0], slider_y[1], num_iterations=iterations)
        avg_diff_2nd_0 = avg_diff_2nd[0].to(torch.float16)
        avg_diff_2nd_1 = avg_diff_2nd[1].to(torch.float16)
        y_concept_1, y_concept_2 = slider_y[0], slider_y[1]
    end_time = time.time()
    print(f"direction time: {end_time - start_time:.2f} ms")
    
    start_time = time.time()
    
    if img2img_type=="controlnet canny" and img is not None:
        control_img = process_controlnet_img(img)
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, image=control_img, controlnet_conditioning_scale =controlnet_scale, scale=0, scale_2nd=0, seed=seed, num_inference_steps=steps, avg_diff=(avg_diff_0,avg_diff_1), avg_diff_2nd=(avg_diff_2nd_0,avg_diff_2nd_1))
    elif img2img_type=="ip adapter" and img is not None:
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, ip_adapter_image=img, scale=0, scale_2nd=0, seed=seed, num_inference_steps=steps, avg_diff=(avg_diff_0,avg_diff_1), avg_diff_2nd=(avg_diff_2nd_0,avg_diff_2nd_1))
    elif img2img_type=="inversion":
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, scale=0, scale_2nd=0, seed=seed, num_inference_steps=steps, avg_diff=(avg_diff_0,avg_diff_1), avg_diff_2nd=(avg_diff_2nd_0,avg_diff_2nd_1), init_latents = init_latents, zs=zs, edit_threshold=edit_threshold, edit_guidance_scale = edit_guidance_scale)
    else: # text to image
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, scale=0, scale_2nd=0, seed=seed, num_inference_steps=steps, avg_diff=(avg_diff_0,avg_diff_1), avg_diff_2nd=(avg_diff_2nd_0,avg_diff_2nd_1))
    
    end_time = time.time()
    print(f"generation time: {end_time - start_time:.2f} ms")
    
    comma_concepts_x = ', '.join(slider_x)
    comma_concepts_y = ', '.join(slider_y)

    avg_diff_x_1 = avg_diff_0.cpu()
    avg_diff_x_2 = avg_diff_1.cpu()
    avg_diff_y_1 = avg_diff_2nd_0.cpu()
    avg_diff_y_2 = avg_diff_2nd_1.cpu()
  
    return gr.update(label=comma_concepts_x, interactive=True),gr.update(label=comma_concepts_y, interactive=True), x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, image

@spaces.GPU
def update_scales(x,y,prompt,seed, steps, guidance_scale,
                  avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, 
                  img2img_type = None, img = None,
                  controlnet_scale= None, ip_adapter_scale=None,
                  edit_threshold=None, edit_guidance_scale = None,
                  init_latents=None, zs=None):
    avg_diff = (avg_diff_x_1.cuda(), avg_diff_x_2.cuda())
    avg_diff_2nd = (avg_diff_y_1.cuda(), avg_diff_y_2.cuda())
    if img2img_type=="controlnet canny" and img is not None:
        control_img = process_controlnet_img(img)
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, image=control_img, controlnet_conditioning_scale =controlnet_scale, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
    elif img2img_type=="ip adapter" and img is not None:
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, ip_adapter_image=img, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
    elif img2img_type=="inversion":
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=(avg_diff_0,avg_diff_1), avg_diff_2nd=(avg_diff_2nd_0,avg_diff_2nd_1),  edit_threshold=edit_threshold, edit_guidance_scale = edit_guidance_scale, init_latents = init_latents, zs=zs)
    else:     
        image = clip_slider.generate(prompt, guidance_scale=guidance_scale, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
    return image

@spaces.GPU
def update_x(x,y,prompt,seed, steps, 
             avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2,
             img2img_type = None,
             img = None):
    avg_diff = (avg_diff_x_1.cuda(), avg_diff_x_2.cuda())
    avg_diff_2nd = (avg_diff_y_1.cuda(), avg_diff_y_2.cuda())
    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
    return image

@spaces.GPU
def update_y(x,y,prompt, seed, steps, 
            avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2,
            img2img_type = None,
            img = None):
    avg_diff = (avg_diff_x_1.cuda(), avg_diff_x_2.cuda())
    avg_diff_2nd = (avg_diff_y_1.cuda(), avg_diff_y_2.cuda())
    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, seed=seed, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
    return image

@spaces.GPU
def invert(image, num_inversion_steps=50, skip=0.3):
    _ = clip_slider_inv.pipe.invert(
    source_prompt = "", 
    image = image,
    num_inversion_steps = num_inversion_steps,
    skip = skip
)
    return clip_slider_inv.pipe.init_latents, lip_slider_inv.pipe.zs

def reset_do_inversion():
        return True
css = '''
#group {
    position: relative;
    width: 420px;
    height: 420px;
    margin-bottom: 20px;
    background-color: white
}
#x {
    position: absolute;
    bottom: 0;
    left: 25px;
    width: 400px;
}
#y {
    position: absolute;
    bottom: 20px;
    left: 67px;
    width: 400px;
    transform: rotate(-90deg);
    transform-origin: left bottom;
}
#image_out{position:absolute; width: 80%; right: 10px; top: 40px}
'''
with gr.Blocks(css=css) as demo:
    
    x_concept_1 = gr.State("")
    x_concept_2 = gr.State("")
    y_concept_1 = gr.State("")
    y_concept_2 = gr.State("")

    avg_diff_x_1 = gr.State()
    avg_diff_x_2 = gr.State()
    avg_diff_y_1 = gr.State()
    avg_diff_y_2 = gr.State()

    do_inversion = gr.State()
    init_latents = gr.State()
    zs = gr.State()
    
    with gr.Tab("text2image"):
        with gr.Row():
            with gr.Column():
                slider_x = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                slider_y = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                prompt = gr.Textbox(label="Prompt")
                submit = gr.Button("find directions")
            with gr.Column():
                with gr.Group(elem_id="group"):
                  x = gr.Slider(minimum=-7, value=0, maximum=7, elem_id="x", interactive=False)
                  y = gr.Slider(minimum=-7, value=0, maximum=7, elem_id="y", interactive=False)
                  output_image = gr.Image(elem_id="image_out")
                with gr.Row():
                    generate_butt = gr.Button("generate")
        
        with gr.Accordion(label="advanced options", open=False):
            iterations = gr.Slider(label = "num iterations", minimum=0, value=200, maximum=400)
            steps = gr.Slider(label = "num inference steps", minimum=1, value=8, maximum=30)
            guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5,
                )
            seed  = gr.Slider(minimum=0, maximum=np.iinfo(np.int32).max, label="Seed", interactive=True, randomize=True)
        
       
    with gr.Tab(label="image2image"):
        with gr.Row():
            with gr.Column():
                image = gr.ImageEditor(type="pil", image_mode="L", crop_size=(512, 512))
                slider_x_a = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                slider_y_a = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                img2img_type = gr.Radio(["controlnet canny", "ip adapter"], label="", info="")
                prompt_a = gr.Textbox(label="Prompt")
                submit_a = gr.Button("Submit")
            with gr.Column():
                with gr.Group(elem_id="group"):
                  x_a = gr.Slider(minimum=-10, value=0, maximum=10, elem_id="x", interactive=False)
                  y_a = gr.Slider(minimum=-10, value=0, maximum=10, elem_id="y", interactive=False)
                  output_image_a = gr.Image(elem_id="image_out")
                with gr.Row():
                    generate_butt_a = gr.Button("generate")
        
        with gr.Accordion(label="advanced options", open=False):
            iterations_a = gr.Slider(label = "num iterations", minimum=0, value=200, maximum=300)
            steps_a = gr.Slider(label = "num inference steps", minimum=1, value=8, maximum=30)
            guidance_scale_a = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5,
                )
            controlnet_conditioning_scale = gr.Slider(
                    label="controlnet conditioning scale",
                    minimum=0.5,
                    maximum=5.0,
                    step=0.1,
                    value=0.7,
                )
            ip_adapter_scale = gr.Slider(
                    label="ip adapter scale",
                    minimum=0.5,
                    maximum=5.0,
                    step=0.1,
                    value=0.8,
                )
            seed_a  = gr.Slider(minimum=0, maximum=np.iinfo(np.int32).max, label="Seed", interactive=True, randomize=True)

    with gr.Tab(label="inversion"):
        with gr.Row():
            with gr.Column():
                image_inv = gr.Image(height=512, width=512)
                slider_x_inv = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                slider_y_inv = gr.Dropdown(label="Slider X concept range", allow_custom_value=True, multiselect=True, max_choices=2)
                prompt_inv = gr.Textbox(label="Prompt")
                img2img_type_inv = gr.Radio(["inversion"], label="",value="inversion", info="", visible=False)
                submit_inv = gr.Button("Submit")
            with gr.Column():
                with gr.Group(elem_id="group"):
                  x_inv = gr.Slider(minimum=-10, value=0, maximum=10, elem_id="x", interactive=False)
                  y_inv = gr.Slider(minimum=-10, value=0, maximum=10, elem_id="y", interactive=False)
                  output_image_inv = gr.Image(elem_id="image_out")
                with gr.Row():
                    generate_butt_inv = gr.Button("generate")
        
        with gr.Accordion(label="advanced options", open=False):
            iterations_inv = gr.Slider(label = "num iterations", minimum=0, value=200, maximum=300)
            steps_inv = gr.Slider(label = "num inference steps", minimum=1, value=8, maximum=30)
            guidance_scale_inv = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5,
                )
            # edit_threshold=None, edit_guidance_scale = None,
            #  init_latents=None, zs=None
            edit_threshold = gr.Slider(
                    label="edit threshold",
                    minimum=0.01,
                    maximum=0.99,
                    step=0.1,
                    value=0.3,
                )
            edit_guidance_scale = gr.Slider(
                    label="edit guidance scale",
                    minimum=0,
                    maximum=20,
                    step=0.25,
                    value=5,
                )
            seed_inv = gr.Slider(minimum=0, maximum=np.iinfo(np.int32).max, label="Seed", interactive=True, randomize=True)
        
    submit.click(fn=generate,
                     inputs=[slider_x, slider_y, prompt, seed, iterations, steps, guidance_scale, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2],
                     outputs=[x, y, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, output_image])

    image_inv.change(fn=reset_do_inversion, outputs=[do_inversion]).then(fn=invert, inputs=[image_inv], outputs=[init_latents,zs])
    submit_inv.click(fn=generate,
                     inputs=[slider_x_inv, slider_y_inv, prompt_inv, seed_inv, iterations_inv, steps_inv, guidance_scale_inv, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type_inv, image, controlnet_conditioning_scale, ip_adapter_scale ,edit_threshold, edit_guidance_scale, init_latents, zs],
                     outputs=[x_inv, y_inv, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, output_image_inv])
    
    generate_butt.click(fn=update_scales, inputs=[x,y, prompt, seed, steps, guidance_scale, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2], outputs=[output_image])
    generate_butt_a.click(fn=update_scales, inputs=[x_a,y_a, prompt_a, seed_a, steps_a, guidance_scale_a, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type, image, controlnet_conditioning_scale, ip_adapter_scale], outputs=[output_image_a])
    generate_butt_inv.click(fn=update_scales, inputs=[x_inv,y_inv, prompt_inv, seed_inv, steps_inv, guidance_scale_inv, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type_inv, image, controlnet_conditioning_scale, ip_adapter_scale ,edit_threshold, edit_guidance_scale, init_latents, zs], outputs=[output_image_inv])
    #x.change(fn=update_scales, inputs=[x,y, prompt, seed, steps, guidance_scale, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2], outputs=[output_image])
    #y.change(fn=update_scales, inputs=[x,y, prompt, seed, steps, guidance_scale, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2], outputs=[output_image])
    submit_a.click(fn=generate,
                     inputs=[slider_x_a, slider_y_a, prompt_a, seed_a, iterations_a, steps_a, guidance_scale_a, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type, image, controlnet_conditioning_scale, ip_adapter_scale],
                     outputs=[x_a, y_a, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, output_image_a])
    #x_a.change(fn=update_scales, inputs=[x_a,y_a, prompt_a, seed_a, steps_a, guidance_scale_a, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type, image, controlnet_conditioning_scale, ip_adapter_scale], outputs=[output_image_a])
    #y_a.change(fn=update_scales, inputs=[x_a,y_a, prompt, seed_a, steps_a, guidance_scale_a, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, img2img_type, image, controlnet_conditioning_scale, ip_adapter_scale], outputs=[output_image_a])

        
if __name__ == "__main__":
    demo.launch()