import os import types from typing import Any, List, Optional import gradio as gr import matplotlib.pyplot as plt import numpy as np import spaces import torch import torch.nn.functional as F from diffusers import StableDiffusion3Pipeline from diffusers.models.attention_processor import Attention import logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) pipe = StableDiffusion3Pipeline.from_pretrained( "stabilityai/stable-diffusion-3-medium-diffusers", text_encoder_3=None, tokenizer_3=None, torch_dtype=torch.bfloat16, ) def encode_concepts( pipeline, concepts: list[str] = None, num_images_per_prompt: int = 1, clip_skip: Optional[int] = None, max_sequence_length: int = 256, device: Optional[torch.device] = None, ): device = device or pipeline._execution_device prompt = " ".join(concepts) prompt = [prompt] if isinstance(prompt, str) else prompt prompt_embed, _ = pipeline._get_clip_prompt_embeds( prompt=prompt, device=device, num_images_per_prompt=num_images_per_prompt, clip_skip=clip_skip, clip_model_index=0, ) prompt_2_embed, _ = pipeline._get_clip_prompt_embeds( prompt=prompt, device=device, num_images_per_prompt=num_images_per_prompt, clip_skip=clip_skip, clip_model_index=1, ) clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) t5_prompt_embed = pipeline._get_t5_prompt_embeds( prompt=prompt, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, device=device, ) clip_prompt_embeds = torch.nn.functional.pad( clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]), ) prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) prompt_embeds = prompt_embeds[:, : len(concepts) + 1] # Leave room for padding token return prompt_embeds def encode_concepts_prompts(concepts, prompt, pipe): concept_embeddings = encode_concepts( pipeline=pipe, concepts=concepts, device=pipe.device, num_images_per_prompt=1 ) prompt_out_embeds = pipe.encode_prompt( prompt=prompt, prompt_2=prompt, prompt_3=prompt, device=pipe.device, num_images_per_prompt=1, do_classifier_free_guidance=True, ) return concept_embeddings, prompt_out_embeds class CustomAttnProcessor: def __init__(self): self.store_out = { 'concept_out': [], 'image_out': [], } if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("JointAttnProcessor2_0 requires PyTorch 2.0 or newer.") def __call__( self, attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, attention_mask: torch.FloatTensor | None = None, concept_hidden_states: torch.FloatTensor | None = None, *args, **kwargs, ) -> torch.FloatTensor: residual = hidden_states batch_size = hidden_states.shape[0] query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) inner_dim = key.shape[-1] head_dim = inner_dim // attn.heads img_q = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) img_k = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) img_v = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) if attn.norm_q is not None: img_q = attn.norm_q(img_q) if attn.norm_k is not None: img_k = attn.norm_k(img_k) if encoder_hidden_states is not None: text_q = attn.add_q_proj(encoder_hidden_states) text_k = attn.add_k_proj(encoder_hidden_states) text_v = attn.add_v_proj(encoder_hidden_states) text_q = text_q.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) text_k = text_k.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) text_v = text_v.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) if attn.norm_added_q is not None: text_q = attn.norm_added_q(text_q) if attn.norm_added_k is not None: text_k = attn.norm_added_k(text_k) out_c = None if concept_hidden_states is not None: concept_q = attn.add_q_proj(concept_hidden_states) concept_k = attn.add_k_proj(concept_hidden_states) concept_v = attn.add_v_proj(concept_hidden_states) concept_q = concept_q.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) concept_k = concept_k.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) concept_v = concept_v.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) if attn.norm_added_q is not None: concept_q = attn.norm_added_q(concept_q) if attn.norm_added_k is not None: concept_k = attn.norm_added_k(concept_k) k_xc = torch.cat([img_k, concept_k], dim=2) v_xc = torch.cat([img_v, concept_v], dim=2) out_c = F.scaled_dot_product_attention(concept_q, k_xc, v_xc, dropout_p=0.0, is_causal=False) out_c = out_c.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) out_c = out_c.to(concept_q.dtype) self.store_out['concept_out'].append(out_c[1].detach().cpu().half()) if encoder_hidden_states is not None: query = torch.cat([img_q, text_q], dim=2) key = torch.cat([img_k, text_k], dim=2) value = torch.cat([img_v, text_v], dim=2) hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) if encoder_hidden_states is not None: hidden_states, encoder_hidden_states = ( hidden_states[:, : residual.shape[1]], hidden_states[:, residual.shape[1] :], ) self.store_out['image_out'].append(hidden_states[1].detach().cpu().half()) if not attn.context_pre_only: encoder_hidden_states = attn.to_add_out(encoder_hidden_states) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) if concept_hidden_states is not None: return hidden_states, encoder_hidden_states, out_c elif encoder_hidden_states is not None: return hidden_states, encoder_hidden_states else: return hidden_states def transformer_block_forward_patch( self, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, joint_attention_kwargs: dict[str, Any] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} concept_hidden_states = joint_attention_kwargs.get("concept_hidden_states", None) if self.use_dual_attention: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( hidden_states, emb=temb ) else: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) if self.context_pre_only: norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) else: norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( encoder_hidden_states, emb=temb ) if concept_hidden_states is not None: if self.context_pre_only: norm_concept_hidden_states = self.norm1_context(concept_hidden_states) else: norm_concept_hidden_states, _, _, _, _ = self.norm1_context( concept_hidden_states, emb=temb ) joint_attention_kwargs["concept_hidden_states"] = norm_concept_hidden_states if concept_hidden_states is not None: attn_output, context_attn_output, concept_attn_output = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, **joint_attention_kwargs, ) else: attn_output, context_attn_output = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, **joint_attention_kwargs, ) attn_output = gate_msa.unsqueeze(1) * attn_output hidden_states = hidden_states + attn_output if self.use_dual_attention: attn_output2 = self.attn2(hidden_states=norm_hidden_states2, **joint_attention_kwargs) attn_output2 = gate_msa2.unsqueeze(1) * attn_output2 hidden_states = hidden_states + attn_output2 norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] ff_output = self.ff(norm_hidden_states) ff_output = gate_mlp.unsqueeze(1) * ff_output hidden_states = hidden_states + ff_output if self.context_pre_only: encoder_hidden_states = None else: context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output if concept_hidden_states is not None and not self.context_pre_only: concept_attn_output = c_gate_msa.unsqueeze(1) * concept_attn_output concept_hidden_states = concept_hidden_states + concept_attn_output norm_concept_hidden_states = self.norm2_context(concept_hidden_states) norm_concept_hidden_states = norm_concept_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] concept_ff_output = self.ff_context(norm_concept_hidden_states) concept_hidden_states = concept_hidden_states + c_gate_mlp.unsqueeze(1) * concept_ff_output joint_attention_kwargs["concept_hidden_states"] = concept_hidden_states return encoder_hidden_states, hidden_states def apply_patched_fn(pipe, custom_processor): for i in range(23): block = pipe.transformer.transformer_blocks[i] block.attn.processor = custom_processor block.forward = types.MethodType(transformer_block_forward_patch, block) def image_gen(pipe, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, concept_embeds, num_steps, img_size, custom_processor): concept_embeds = pipe.transformer.context_embedder(concept_embeds) output = pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, num_inference_steps=num_steps, guidance_scale=7.0, height=img_size, width=img_size, joint_attention_kwargs={"concept_hidden_states": concept_embeds} ) concept_vector = torch.stack(custom_processor.store_out['concept_out']).view(num_steps, 23, 1, -1, 1536) image_vector = torch.stack(custom_processor.store_out['image_out']).view(num_steps, 23, 1, -1, 1536) return concept_vector, image_vector, output.images[0] def compute_heatmaps( img_info, concept_vectors: torch.Tensor, image_vectors: torch.Tensor, height: int, width: int, layer_indices: Optional[List[int]] = None, softmax: bool = True, average_over_timesteps: bool = False, ) -> torch.Tensor: if average_over_timesteps and len(concept_vectors.shape) == 5: concept_vectors = torch.mean(concept_vectors, dim=0) image_vectors = torch.mean(image_vectors, dim=0) num_layers = concept_vectors.shape[0] if layer_indices is None: layer_indices = list(range(num_layers)) concept_vectors = concept_vectors[:, :, 1:] concept_vectors = concept_vectors / (concept_vectors.norm(dim=-1, keepdim=True) + 1e-8) heatmaps = torch.einsum( "lbcd,lbpd->lbcp", concept_vectors.float(), image_vectors.float() ) if softmax: heatmaps = torch.nn.functional.softmax(heatmaps, dim=-2) heatmaps = heatmaps[layer_indices] heatmaps = heatmaps.mean(dim=0) heatmaps = heatmaps[0] h_tokens = height // (img_info['vae_scale_factor'] * img_info['patch_size']) w_tokens = width // (img_info['vae_scale_factor'] * img_info['patch_size']) heatmaps = heatmaps.view(-1, h_tokens, w_tokens) return heatmaps custom_processor = CustomAttnProcessor() apply_patched_fn(pipe, custom_processor) @spaces.GPU(duration=60) def inference_pipeline(prompt, concepts_string, steps, img_size): try: device = "cuda" if torch.cuda.is_available() else "cpu" pipe.to(device) custom_processor.store_out['concept_out'].clear() custom_processor.store_out['image_out'].clear() IMAGE_SIZE = int(img_size) concepts = [c.strip() for c in concepts_string.split(" ") if c.strip()] with torch.no_grad(): logging.info("Encoding concept and prompt for embeddings") concept_embeddings, ( prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, ) = encode_concepts_prompts(concepts, prompt, pipe) # print(concept_embeddings.shape) logging.info("Generating Image") concept_vector, image_vector, generated_image = image_gen( pipe=pipe, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, concept_embeds=concept_embeddings, num_steps=int(steps), img_size=IMAGE_SIZE, custom_processor=custom_processor ) # print(concept_vector.shape, image_vector.shape) img_info = { "vae_scale_factor": pipe.vae_scale_factor, "patch_size": pipe.patch_size, } # del pipe logging.info("Creating heatmaps") heatmaps = compute_heatmaps( img_info, concept_vectors=concept_vector, image_vectors=image_vector, height=IMAGE_SIZE, width=IMAGE_SIZE, average_over_timesteps=True, ) cols = (len(concepts) + 2) // 2 fig, axes = plt.subplots(2, cols, figsize=(12, 6)) axes_flat = axes.flatten() img = np.array(generated_image) for idx in range(len(concepts) + 1): if idx == 0: axes_flat[idx].imshow(img) axes_flat[idx].set_title("Original Image") else: hm = heatmaps[idx - 1].detach().cpu().numpy() hm = (hm - hm.min()) / (hm.max() - hm.min() + 1e-8) axes_flat[idx].imshow(hm, cmap="plasma") axes_flat[idx].set_title(concepts[idx - 1]) axes_flat[idx].axis("off") for idx in range(len(concepts) + 1, len(axes_flat)): axes_flat[idx].axis("off") plt.tight_layout() return fig except Exception as e: import traceback traceback.print_exc() raise gr.Error(f"Pipeline Failed: {str(e)}") finally: import gc concept_vector = None image_vector = None output = None gc.collect() # torch.cuda.empty_cache() ui = gr.Interface( fn=inference_pipeline, inputs=[ gr.Textbox(label="Image Prompt", value="A dragon standing on a rock."), gr.Textbox(label="Concepts (Space-Separated)", value="dragon rock sky"), gr.Slider(minimum=1, maximum=30, value=28, step=1, label="SD3 Inference Steps"), gr.Slider(minimum=512, maximum=1024, value=768, step=64, label="Output Image Size"), ], outputs=gr.Plot(label="Concept Attention Maps"), title="ConceptAttention Visualization", ) if __name__ == "__main__": ui.launch()