|
import os |
|
from typing import Optional |
|
|
|
import gradio as gr |
|
import numpy as np |
|
import pandas as pd |
|
import torch |
|
from PIL import Image |
|
from scipy.stats import beta as beta_distribution |
|
|
|
from pipeline_interpolated_sdxl import InterpolationStableDiffusionXLPipeline |
|
from pipeline_interpolated_stable_diffusion import InterpolationStableDiffusionPipeline |
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
|
title = r""" |
|
<h1 align="center">PAID: (Prompt-guided) Attention Interpolation of Text-to-Image Diffusion</h1> |
|
""" |
|
|
|
description = r""" |
|
<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/QY-H00/attention-interpolation-diffusion/tree/public' target='_blank'><b>PAID: (Prompt-guided) Attention Interpolation of Text-to-Image Diffusion</b></a>.<br> |
|
How to use:<br> |
|
1. Input prompt 1 and prompt 2. |
|
2. (Optional) Input the guidance prompt and negative prompt. |
|
3. (Optional) Change the interpolation parameters and check the Beta distribution. |
|
4. Click the <b>Generate</b> button to begin generating images. |
|
5. Enjoy! 😊""" |
|
|
|
article = r""" |
|
--- |
|
✒️ **Citation** |
|
<br> |
|
If you found this demo/our paper useful, please consider citing: |
|
```bibtex |
|
@article{he024paid, |
|
title={PAID:(Prompt-guided) Attention Interpolation of Text-to-Image Diffusion}, |
|
author={He, Qiyuan and Wang, Jinghao and Liu, Ziwei and Angle, Yao}, |
|
journal={}, |
|
year={2024} |
|
} |
|
``` |
|
📧 **Contact** |
|
<br> |
|
If you have any questions, please feel free to open an issue in our <a href='https://github.com/QY-H00/attention-interpolation-diffusion/tree/public' target='_blank'><b>Github Repo</b></a> or directly reach us out at <b>qhe@u.nus.edu.sg</b>. |
|
""" |
|
|
|
MAX_SEED = np.iinfo(np.int32).max |
|
CACHE_EXAMPLES = False |
|
USE_TORCH_COMPILE = False |
|
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1" |
|
PREVIEW_IMAGES = False |
|
|
|
dtype = torch.float32 |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
pipeline = InterpolationStableDiffusionPipeline( |
|
repo_name="runwayml/stable-diffusion-v1-5", |
|
guidance_scale=10.0, |
|
scheduler_name="unipc", |
|
) |
|
pipeline.to(device, dtype=dtype) |
|
|
|
|
|
def change_model_fn(model_name: str) -> None: |
|
global pipeline |
|
name_mapping = { |
|
"SD1.4-521": "CompVis/stable-diffusion-v1-4", |
|
"SD1.5-512": "runwayml/stable-diffusion-v1-5", |
|
"SD2.1-768": "stabilityai/stable-diffusion-2-1", |
|
"SDXL-1024": "stabilityai/stable-diffusion-xl-base-1.0", |
|
} |
|
if "XL" not in model_name: |
|
pipeline = InterpolationStableDiffusionPipeline( |
|
repo_name=name_mapping[model_name], |
|
guidance_scale=10.0, |
|
scheduler_name="unipc", |
|
) |
|
pipeline.to(device, dtype=dtype) |
|
else: |
|
pipeline = InterpolationStableDiffusionXLPipeline.from_pretrained( |
|
name_mapping[model_name] |
|
) |
|
pipeline.to(device, dtype=dtype) |
|
|
|
|
|
def save_image(img, index): |
|
unique_name = f"{index}.png" |
|
img = Image.fromarray(img) |
|
img.save(unique_name) |
|
return unique_name |
|
|
|
|
|
def generate_beta_tensor( |
|
size: int, alpha: float = 3.0, beta: float = 3.0 |
|
) -> torch.FloatTensor: |
|
prob_values = [i / (size - 1) for i in range(size)] |
|
inverse_cdf_values = beta_distribution.ppf(prob_values, alpha, beta) |
|
return inverse_cdf_values |
|
|
|
|
|
def plot_gemma_fn(alpha: float, beta: float, size: int) -> pd.DataFrame: |
|
beta_ppf = generate_beta_tensor(size=size, alpha=int(alpha), beta=int(beta)) |
|
return pd.DataFrame( |
|
{ |
|
"interpolation index": [i for i in range(size)], |
|
"coefficient": beta_ppf.tolist(), |
|
} |
|
) |
|
|
|
|
|
def get_example() -> list: |
|
case = [ |
|
[ |
|
"A photo of dog, best quality, extremely detailed", |
|
"A photo of car, best quality, extremely detailed", |
|
3, |
|
6, |
|
3, |
|
"A photo of a dog driving a car, logical, best quality, extremely detailed", |
|
"monochrome, lowres, bad anatomy, worst quality, low quality", |
|
"SD1.5-512", |
|
6.1 / 50, |
|
10, |
|
50, |
|
"fused_inner", |
|
"self", |
|
1002, |
|
True, |
|
] |
|
] |
|
return case |
|
|
|
|
|
def dynamic_gallery_fn(interpolation_size: int): |
|
|
|
return gr.Gallery( |
|
label="Result", show_label=False, rows=1, columns=interpolation_size |
|
) |
|
|
|
|
|
@torch.no_grad() |
|
def generate( |
|
prompt1: str, |
|
prompt2: str, |
|
guidance_prompt: Optional[str] = None, |
|
negative_prompt: str = "", |
|
warmup_ratio: int = 8, |
|
guidance_scale: float = 10, |
|
early: str = "fused_outer", |
|
late: str = "self", |
|
alpha: float = 4.0, |
|
beta: float = 4.0, |
|
interpolation_size: int = 3, |
|
seed: int = 0, |
|
same_latent: bool = True, |
|
num_inference_steps: int = 50, |
|
progress=gr.Progress(), |
|
) -> np.ndarray: |
|
global pipeline |
|
generator = ( |
|
torch.cuda.manual_seed(seed) |
|
if torch.cuda.is_available() |
|
else torch.manual_seed(seed) |
|
) |
|
latent1 = pipeline.generate_latent(generator=generator) |
|
latent1 = latent1.to(device=pipeline.unet.device, dtype=pipeline.unet.dtype) |
|
if same_latent: |
|
latent2 = latent1.clone() |
|
else: |
|
latent2 = pipeline.generate_latent(generator=generator) |
|
latent2 = latent2.to(device=pipeline.unet.device, dtype=pipeline.unet.dtype) |
|
betas = generate_beta_tensor(size=interpolation_size, alpha=alpha, beta=beta) |
|
for i in progress.tqdm( |
|
range(interpolation_size - 2), |
|
desc=( |
|
f"Generating {interpolation_size-2} images" |
|
if interpolation_size > 3 |
|
else "Generating 1 image" |
|
), |
|
): |
|
it = betas[i + 1].item() |
|
images = pipeline.interpolate_single( |
|
it, |
|
latent_start=latent1, |
|
latent_end=latent2, |
|
prompt_start=prompt1, |
|
prompt_end=prompt2, |
|
guide_prompt=guidance_prompt, |
|
num_inference_steps=num_inference_steps, |
|
warmup_ratio=warmup_ratio, |
|
early=early, |
|
late=late, |
|
negative_prompt=negative_prompt, |
|
guidance_scale=guidance_scale, |
|
) |
|
if interpolation_size == 3: |
|
final_images = images |
|
break |
|
if i == 0: |
|
final_images = images[:2] |
|
elif i == interpolation_size - 3: |
|
final_images = np.concatenate([final_images, images[1:]], axis=0) |
|
else: |
|
final_images = np.concatenate([final_images, images[1:2]], axis=0) |
|
return final_images |
|
|
|
|
|
interpolation_size = None |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown(title) |
|
gr.Markdown(description) |
|
with gr.Group(): |
|
prompt1 = gr.Text( |
|
label="Prompt 1", |
|
max_lines=3, |
|
placeholder="Enter the First Prompt", |
|
interactive=True, |
|
value="A photo of dog, best quality, extremely detailed", |
|
) |
|
prompt2 = gr.Text( |
|
label="Prompt 2", |
|
max_lines=3, |
|
placeholder="Enter the Second prompt", |
|
interactive=True, |
|
value="A photo of car, best quality, extremely detaile", |
|
) |
|
result = gr.Gallery(label="Result", show_label=False, rows=1, columns=3) |
|
generate_button = gr.Button("Generate", variant="primary") |
|
with gr.Accordion("Advanced options", open=True): |
|
with gr.Group(): |
|
with gr.Row(): |
|
with gr.Column(): |
|
interpolation_size = gr.Slider( |
|
label="Interpolation Size", |
|
minimum=3, |
|
maximum=15, |
|
step=1, |
|
value=3, |
|
info="Interpolation size includes the start and end images", |
|
) |
|
alpha = gr.Slider( |
|
label="alpha", |
|
minimum=1, |
|
maximum=50, |
|
step=0.1, |
|
value=6.0, |
|
) |
|
beta = gr.Slider( |
|
label="beta", |
|
minimum=1, |
|
maximum=50, |
|
step=0.1, |
|
value=3.0, |
|
) |
|
gamma_plot = gr.LinePlot( |
|
x="interpolation index", |
|
y="coefficient", |
|
title="Beta Distribution with Sampled Points", |
|
height=500, |
|
width=400, |
|
overlay_point=True, |
|
tooltip=["coefficient", "interpolation index"], |
|
interactive=False, |
|
show_label=False, |
|
) |
|
gamma_plot.change( |
|
plot_gemma_fn, |
|
inputs=[ |
|
alpha, |
|
beta, |
|
interpolation_size, |
|
], |
|
outputs=gamma_plot, |
|
) |
|
with gr.Group(): |
|
guidance_prompt = gr.Text( |
|
label="Guidance prompt", |
|
max_lines=3, |
|
placeholder="Enter a Guidance Prompt", |
|
interactive=True, |
|
value="A photo of a dog driving a car, logical, best quality, extremely detailed", |
|
) |
|
negative_prompt = gr.Text( |
|
label="Negative prompt", |
|
max_lines=3, |
|
placeholder="Enter a Negative Prompt", |
|
interactive=True, |
|
value="monochrome, lowres, bad anatomy, worst quality, low quality", |
|
) |
|
with gr.Row(): |
|
with gr.Column(): |
|
warmup_ratio = gr.Slider( |
|
label="Warmup Ratio", |
|
minimum=0.02, |
|
maximum=1, |
|
step=0.01, |
|
value=0.122, |
|
interactive=True, |
|
) |
|
guidance_scale = gr.Slider( |
|
label="Guidance Scale", |
|
minimum=0, |
|
maximum=50, |
|
step=0.1, |
|
value=10, |
|
interactive=True, |
|
) |
|
with gr.Column(): |
|
early = gr.Dropdown( |
|
label="Early stage attention type", |
|
choices=[ |
|
"pure_inner", |
|
"fused_inner", |
|
"pure_outer", |
|
"fused_outer", |
|
"self", |
|
], |
|
value="fused_outer", |
|
type="value", |
|
interactive=True, |
|
) |
|
late = gr.Dropdown( |
|
label="Late stage attention type", |
|
choices=[ |
|
"pure_inner", |
|
"fused_inner", |
|
"pure_outer", |
|
"fused_outer", |
|
"self", |
|
], |
|
value="self", |
|
type="value", |
|
interactive=True, |
|
) |
|
num_inference_steps = gr.Slider( |
|
label="Inference Steps", |
|
minimum=25, |
|
maximum=50, |
|
step=1, |
|
value=50, |
|
interactive=True, |
|
) |
|
with gr.Row(): |
|
model_choice = gr.Dropdown( |
|
["SD1.4-521", "SD1.5-512", "SD2.1-768", "SDXL-1024"], |
|
label="Model", |
|
value="SD1.5-512", |
|
interactive=True, |
|
) |
|
with gr.Column(): |
|
seed = gr.Slider( |
|
label="Seed", |
|
minimum=0, |
|
maximum=MAX_SEED, |
|
step=1, |
|
value=1002, |
|
) |
|
same_latent = gr.Checkbox( |
|
label="Same latent", |
|
value=True, |
|
info="Use the same latent for start and end images", |
|
show_label=True, |
|
) |
|
|
|
gr.Examples( |
|
examples=get_example(), |
|
inputs=[ |
|
prompt1, |
|
prompt2, |
|
interpolation_size, |
|
alpha, |
|
beta, |
|
guidance_prompt, |
|
negative_prompt, |
|
model_choice, |
|
warmup_ratio, |
|
guidance_scale, |
|
num_inference_steps, |
|
early, |
|
late, |
|
seed, |
|
same_latent, |
|
], |
|
outputs=result, |
|
fn=generate, |
|
cache_examples=CACHE_EXAMPLES, |
|
) |
|
|
|
alpha.change( |
|
fn=plot_gemma_fn, inputs=[alpha, beta, interpolation_size], outputs=gamma_plot |
|
) |
|
beta.change( |
|
fn=plot_gemma_fn, inputs=[alpha, beta, interpolation_size], outputs=gamma_plot |
|
) |
|
interpolation_size.change( |
|
fn=plot_gemma_fn, inputs=[alpha, beta, interpolation_size], outputs=gamma_plot |
|
) |
|
model_choice.change(fn=change_model_fn, inputs=model_choice) |
|
inputs = [ |
|
prompt1, |
|
prompt2, |
|
guidance_prompt, |
|
negative_prompt, |
|
warmup_ratio, |
|
guidance_scale, |
|
early, |
|
late, |
|
alpha, |
|
beta, |
|
interpolation_size, |
|
seed, |
|
same_latent, |
|
num_inference_steps, |
|
] |
|
generate_button.click( |
|
fn=dynamic_gallery_fn, |
|
inputs=interpolation_size, |
|
outputs=result, |
|
).then( |
|
fn=generate, |
|
inputs=inputs, |
|
outputs=result, |
|
) |
|
gr.Markdown(article) |
|
|
|
with gr.Blocks(css="style.css") as demo_with_history: |
|
with gr.Tab("App"): |
|
demo.render() |
|
|
|
if __name__ == "__main__": |
|
demo_with_history.queue(max_size=20).launch() |
|
|
