Add SD3

app.py

@@ -1,25 +1,25 @@
 import torch
 import gradio as gr
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusion3Pipeline
 from utils import (
-    cross_attn_init,
-    register_cross_attention_hook,
     attn_maps,
-    get_net_attn_map,
-    resize_net_attn_map,
-    return_net_attn_map,
+    cross_attn_init,
+    init_pipeline,
+    save_attention_maps
 )
 # from transformers.utils.hub import move_cache
 
 # move_cache()
 
 cross_attn_init()
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    # "stabilityai/sdxl-turbo",
-    torch_dtype=torch.float16,
+
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    "stabilityai/stable-diffusion-3-medium-diffusers",
+    torch_dtype=torch.bfloat16
 )
-pipe.unet = register_cross_attention_hook(pipe.unet)
+
+pipe = init_pipeline(pipe)
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 pipe = pipe.to(device)
 
@@ -29,15 +29,10 @@ def inference(prompt):
         prompt,
         num_inference_steps=15,
     ).images[0]
-    net_attn_maps = get_net_attn_map(image.size)
-    net_attn_maps = resize_net_attn_map(net_attn_maps, image.size)
-    net_attn_maps = return_net_attn_map(net_attn_maps, pipe.tokenizer, prompt)
 
-    # remove sos and eos
-    net_attn_maps = [attn_map for attn_map in net_attn_maps if attn_map[1].split('_')[-1] != "<<|startoftext|>>"]
-    net_attn_maps = [attn_map for attn_map in net_attn_maps if attn_map[1].split('_')[-1] != "<<|endoftext|>>"]
+    total_attn_maps = save_attention_maps(attn_maps, tokenizer, prompts)
     
-    return image, net_attn_maps
+    return image, total_attn_maps
 
 
 with gr.Blocks() as demo:
@@ -46,8 +41,7 @@ with gr.Blocks() as demo:
         # 🚀 Text-to-Image Cross Attention Map for 🧨 Diffusers ⚡
         """
     )
-    # prompt = gr.Textbox(value="A photo of a black puppy, christmas atmosphere", label="Prompt", lines=2)
-    prompt = gr.Textbox(value="A portrait photo of a kangaroo wearing an orange hoodie and blue sunglasses standing on the grass in front of the Sydney Opera House holding a sign on the chest that says 'SDXL'!.", label="Prompt", lines=2)
+    prompt = gr.Textbox(value="A capybara holding a sign that reads Hello World.", label="Prompt", lines=2)
     btn = gr.Button("Generate images", scale=0)
     
     with gr.Row():