import torch
from pipelines.inverted_ve_pipeline import CrossFrameAttnProcessor, CrossFrameAttnProcessor_store, ACTIVATE_LAYER_CANDIDATE
from diffusers import DDIMScheduler, AutoencoderKL
import os
from PIL import Image
from utils import memory_efficient
from diffusers.models.attention_processor import AttnProcessor
from pipeline_stable_diffusion_xl_attn import StableDiffusionXLPipeline
def create_image_grid(image_list, rows, cols, padding=10):
# Ensure the number of rows and columns doesn't exceed the number of images
rows = min(rows, len(image_list))
cols = min(cols, len(image_list))
# Get the dimensions of a single image
image_width, image_height = image_list[0].size
# Calculate the size of the output image
grid_width = cols * (image_width + padding) - padding
grid_height = rows * (image_height + padding) - padding
# Create an empty grid image
grid_image = Image.new('RGB', (grid_width, grid_height), (255, 255, 255))
# Paste images into the grid
for i, img in enumerate(image_list[:rows * cols]):
row = i // cols
col = i % cols
x = col * (image_width + padding)
y = row * (image_height + padding)
grid_image.paste(img, (x, y))
return grid_image
def transform_variable_name(input_str, attn_map_save_step):
# Split the input string into parts using the dot as a separator
parts = input_str.split('.')
# Extract numerical indices from the parts
indices = [int(part) if part.isdigit() else part for part in parts]
# Build the desired output string
output_str = f'pipe.unet.{indices[0]}[{indices[1]}].{indices[2]}[{indices[3]}].{indices[4]}[{indices[5]}].{indices[6]}.attn_map[{attn_map_save_step}]'
return output_str
num_images_per_prompt = 4
seeds=[1] #craft_clay
activate_layer_indices_list = [
# ((0,28),(108,140)),
# ((0,48), (68,140)),
# ((0,48), (88,140)),
# ((0,48), (108,140)),
# ((0,48), (128,140)),
# ((0,48), (140,140)),
# ((0,28), (68,140)),
# ((0,28), (88,140)),
# ((0,28), (108,140)),
# ((0,28), (128,140)),
# ((0,28), (140,140)),
# ((0,8), (68,140)),
# ((0,8), (88,140)),
# ((0,8), (108,140)),
# ((0,8), (128,140)),
# ((0,8), (140,140)),
# ((0,0), (68,140)),
# ((0,0), (88,140)),
((0,0), (108,140)),
# ((0,0), (128,140)),
# ((0,0), (140,140))
]
save_layer_list = [
# 'up_blocks.0.attentions.1.transformer_blocks.0.attn1.processor', #68
# 'up_blocks.0.attentions.1.transformer_blocks.4.attn2.processor', #78
# 'up_blocks.0.attentions.2.transformer_blocks.0.attn1.processor', #88
# 'up_blocks.0.attentions.2.transformer_blocks.4.attn2.processor', #108
# 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.processor', #128
# 'up_blocks.1.attentions.2.transformer_blocks.1.attn1.processor', #138
'up_blocks.0.attentions.2.transformer_blocks.0.attn1.processor', #108
'up_blocks.0.attentions.2.transformer_blocks.0.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.1.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.1.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.2.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.2.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.3.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.3.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.4.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.4.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.5.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.5.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.6.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.6.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.7.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.7.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.8.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.8.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.9.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.9.attn2.processor',
'up_blocks.1.attentions.0.transformer_blocks.0.attn1.processor', #128
'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.0.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.0.transformer_blocks.1.attn2.processor',
'up_blocks.1.attentions.1.transformer_blocks.0.attn1.processor',
'up_blocks.1.attentions.1.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.1.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.1.transformer_blocks.1.attn2.processor',
'up_blocks.1.attentions.2.transformer_blocks.0.attn1.processor',
'up_blocks.1.attentions.2.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.2.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.2.transformer_blocks.1.attn2.processor',
]
attn_map_save_steps = [20]
# attn_map_save_steps = [10,20,30,40]
results_dir = 'saved_attention_map_results'
if not os.path.exists(results_dir):
os.makedirs(results_dir)
base_model_path = "runwayml/stable-diffusion-v1-5"
vae_model_path = "stabilityai/sd-vae-ft-mse"
image_encoder_path = "models/image_encoder/"
object_list = [
"cat",
# "woman",
# "dog",
# "horse",
# "motorcycle"
]
target_object_list = [
# "Null",
"dog",
# "clock",
# "car"
# "panda",
# "bridge",
# "flower"
]
prompt_neg_prompt_pair_dicts = {
# "line_art": ("line art drawing {prompt} . professional, sleek, modern, minimalist, graphic, line art, vector graphics",
# "anime, photorealistic, 35mm film, deformed, glitch, blurry, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, mutated, realism, realistic, impressionism, expressionism, oil, acrylic"
# ) ,
# "anime": ("anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed",
# "photo, deformed, black and white, realism, disfigured, low contrast"
# ),
# "Artstyle_Pop_Art" : ("pop Art style {prompt} . bright colors, bold outlines, popular culture themes, ironic or kitsch",
# "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, minimalist"
# ),
# "Artstyle_Pointillism": ("pointillism style {prompt} . composed entirely of small, distinct dots of color, vibrant, highly detailed",
# "line drawing, smooth shading, large color fields, simplistic"
# ),
# "origami": ("origami style {prompt} . paper art, pleated paper, folded, origami art, pleats, cut and fold, centered composition",
# "noisy, sloppy, messy, grainy, highly detailed, ultra textured, photo"
# ),
"craft_clay": ("play-doh style {prompt} . sculpture, clay art, centered composition, Claymation",
"sloppy, messy, grainy, highly detailed, ultra textured, photo"
),
# "low_poly" : ("low-poly style {prompt} . low-poly game art, polygon mesh, jagged, blocky, wireframe edges, centered composition",
# "noisy, sloppy, messy, grainy, highly detailed, ultra textured, photo"
# ),
# "Artstyle_watercolor": ("watercolor painting {prompt} . vibrant, beautiful, painterly, detailed, textural, artistic",
# "anime, photorealistic, 35mm film, deformed, glitch, low contrast, noisy"
# ),
# "Papercraft_Collage" : ("collage style {prompt} . mixed media, layered, textural, detailed, artistic",
# "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic"
# ),
# "Artstyle_Impressionist" : ("impressionist painting {prompt} . loose brushwork, vibrant color, light and shadow play, captures feeling over form",
# "anime, photorealistic, 35mm film, deformed, glitch, low contrast, noisy"
# )
}
noise_scheduler = DDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,
set_alpha_to_one=False,
steps_offset=1,
)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if device == 'cpu':
torch_dtype = torch.float32
else:
torch_dtype = torch.float16
vae = AutoencoderKL.from_pretrained(vae_model_path, torch_dtype=torch_dtype)
pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch_dtype)
memory_efficient(vae, device)
memory_efficient(pipe, device)
for seed in seeds:
for activate_layer_indices in activate_layer_indices_list:
attn_procs = {}
activate_layers = []
str_activate_layer = ""
for activate_layer_index in activate_layer_indices:
activate_layers += ACTIVATE_LAYER_CANDIDATE[activate_layer_index[0]:activate_layer_index[1]]
str_activate_layer += str(activate_layer_index)
for name in pipe.unet.attn_processors.keys():
if name in activate_layers:
if name in save_layer_list:
print(f"layer:{name}")
attn_procs[name] = CrossFrameAttnProcessor_store(unet_chunk_size=2, attn_map_save_steps=attn_map_save_steps)
else:
print(f"layer:{name}")
attn_procs[name] = CrossFrameAttnProcessor(unet_chunk_size=2)
else :
attn_procs[name] = AttnProcessor()
pipe.unet.set_attn_processor(attn_procs)
for target_object in target_object_list:
target_prompt = f"A photo of a {target_object}"
for object in object_list:
for key in prompt_neg_prompt_pair_dicts.keys():
prompt, negative_prompt = prompt_neg_prompt_pair_dicts[key]
generator = torch.Generator(device).manual_seed(seed) if seed is not None else None
images = pipe(
prompt=prompt.replace("{prompt}", object),
guidance_scale = 7.0,
num_images_per_prompt = num_images_per_prompt,
target_prompt = target_prompt,
generator=generator,
)[0]
#make grid
grid = create_image_grid(images, 1, num_images_per_prompt)
save_name = f"{key}_src_{object}_tgt_{target_object}_activate_layer_{str_activate_layer}_seed_{seed}.png"
save_path = os.path.join(results_dir, save_name)
grid.save(save_path)
print("Saved image to: ", save_path)
#save attn map
for attn_map_save_step in attn_map_save_steps:
attn_map_save_name = f"attn_map_raw_{key}_src_{object}_tgt_{target_object}_activate_layer_{str_activate_layer}_attn_map_step_{attn_map_save_step}_seed_{seed}.pt"
attn_map_dic = {}
# for activate_layer in activate_layers:
for activate_layer in save_layer_list:
attn_map_var_name = transform_variable_name(activate_layer, attn_map_save_step)
exec(f"attn_map_dic[\"{activate_layer}\"] = {attn_map_var_name}")
torch.save(attn_map_dic, os.path.join(results_dir, attn_map_save_name))
print("Saved attn map to: ", os.path.join(results_dir, attn_map_save_name))