CrossAttention_Release_NoImages.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b4974c-6437-422d-afae-daa2884ad633",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import CLIPModel, CLIPTextModel, CLIPTokenizer\n",
    "from diffusers import AutoencoderKL, UNet2DConditionModel\n",
    "\n",
    "#NOTE: Last tested working diffusers version is diffusers==0.4.1, https://github.com/huggingface/diffusers/releases/tag/v0.4.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ebfbd7-5026-4830-93e5-d43272db8912",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Init CLIP tokenizer and model\n",
    "model_path_clip = \"openai/clip-vit-large-patch14\"\n",
    "clip_tokenizer = CLIPTokenizer.from_pretrained(model_path_clip)\n",
    "clip_model = CLIPModel.from_pretrained(model_path_clip, torch_dtype=torch.float16)\n",
    "clip = clip_model.text_model\n",
    "\n",
    "#Init diffusion model\n",
    "auth_token = True #Replace this with huggingface auth token as a string if model is not already downloaded\n",
    "model_path_diffusion = \"CompVis/stable-diffusion-v1-4\"\n",
    "unet = UNet2DConditionModel.from_pretrained(model_path_diffusion, subfolder=\"unet\", use_auth_token=auth_token, revision=\"fp16\", torch_dtype=torch.float16)\n",
    "vae = AutoencoderKL.from_pretrained(model_path_diffusion, subfolder=\"vae\", use_auth_token=auth_token, revision=\"fp16\", torch_dtype=torch.float16)\n",
    "\n",
    "#Move to GPU\n",
    "device = \"cuda\"\n",
    "unet.to(device)\n",
    "vae.to(device)\n",
    "clip.to(device)\n",
    "print(\"Loaded all models\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08dee7d1-e050-43d3-86a9-5776276aad78",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from PIL import Image\n",
    "from diffusers import LMSDiscreteScheduler\n",
    "from tqdm.auto import tqdm\n",
    "from torch import autocast\n",
    "from difflib import SequenceMatcher\n",
    "\n",
    "def init_attention_weights(weight_tuples):\n",
    "    tokens_length = clip_tokenizer.model_max_length\n",
    "    weights = torch.ones(tokens_length)\n",
    "    \n",
    "    for i, w in weight_tuples:\n",
    "        if i < tokens_length and i >= 0:\n",
    "            weights[i] = w\n",
    "    \n",
    "    \n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn2\" in name:\n",
    "            module.last_attn_slice_weights = weights.to(device)\n",
    "        if module_name == \"CrossAttention\" and \"attn1\" in name:\n",
    "            module.last_attn_slice_weights = None\n",
    "    \n",
    "\n",
    "def init_attention_edit(tokens, tokens_edit):\n",
    "    tokens_length = clip_tokenizer.model_max_length\n",
    "    mask = torch.zeros(tokens_length)\n",
    "    indices_target = torch.arange(tokens_length, dtype=torch.long)\n",
    "    indices = torch.zeros(tokens_length, dtype=torch.long)\n",
    "\n",
    "    tokens = tokens.input_ids.numpy()[0]\n",
    "    tokens_edit = tokens_edit.input_ids.numpy()[0]\n",
    "    \n",
    "    for name, a0, a1, b0, b1 in SequenceMatcher(None, tokens, tokens_edit).get_opcodes():\n",
    "        if b0 < tokens_length:\n",
    "            if name == \"equal\" or (name == \"replace\" and a1-a0 == b1-b0):\n",
    "                mask[b0:b1] = 1\n",
    "                indices[b0:b1] = indices_target[a0:a1]\n",
    "\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn2\" in name:\n",
    "            module.last_attn_slice_mask = mask.to(device)\n",
    "            module.last_attn_slice_indices = indices.to(device)\n",
    "        if module_name == \"CrossAttention\" and \"attn1\" in name:\n",
    "            module.last_attn_slice_mask = None\n",
    "            module.last_attn_slice_indices = None\n",
    "\n",
    "\n",
    "def init_attention_func():\n",
    "    #ORIGINAL SOURCE CODE: https://github.com/huggingface/diffusers/blob/91ddd2a25b848df0fa1262d4f1cd98c7ccb87750/src/diffusers/models/attention.py#L276\n",
    "    def new_attention(self, query, key, value):\n",
    "        # TODO: use baddbmm for better performance\n",
    "        attention_scores = torch.matmul(query, key.transpose(-1, -2)) * self.scale\n",
    "        attn_slice = attention_scores.softmax(dim=-1)\n",
    "        # compute attention output\n",
    "        \n",
    "        if self.use_last_attn_slice:\n",
    "            if self.last_attn_slice_mask is not None:\n",
    "                new_attn_slice = torch.index_select(self.last_attn_slice, -1, self.last_attn_slice_indices)\n",
    "                attn_slice = attn_slice * (1 - self.last_attn_slice_mask) + new_attn_slice * self.last_attn_slice_mask\n",
    "            else:\n",
    "                attn_slice = self.last_attn_slice\n",
    "\n",
    "            self.use_last_attn_slice = False\n",
    "\n",
    "        if self.save_last_attn_slice:\n",
    "            self.last_attn_slice = attn_slice\n",
    "            self.save_last_attn_slice = False\n",
    "\n",
    "        if self.use_last_attn_weights and self.last_attn_slice_weights is not None:\n",
    "            attn_slice = attn_slice * self.last_attn_slice_weights\n",
    "            self.use_last_attn_weights = False\n",
    "        \n",
    "        hidden_states = torch.matmul(attn_slice, value)\n",
    "        # reshape hidden_states\n",
    "        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)\n",
    "        return hidden_states\n",
    "    \n",
    "    def new_sliced_attention(self, query, key, value, sequence_length, dim):\n",
    "        \n",
    "        batch_size_attention = query.shape[0]\n",
    "        hidden_states = torch.zeros(\n",
    "            (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype\n",
    "        )\n",
    "        slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]\n",
    "        for i in range(hidden_states.shape[0] // slice_size):\n",
    "            start_idx = i * slice_size\n",
    "            end_idx = (i + 1) * slice_size\n",
    "            attn_slice = (\n",
    "                torch.matmul(query[start_idx:end_idx], key[start_idx:end_idx].transpose(1, 2)) * self.scale\n",
    "            )  # TODO: use baddbmm for better performance\n",
    "            attn_slice = attn_slice.softmax(dim=-1)\n",
    "            \n",
    "            if self.use_last_attn_slice:\n",
    "                if self.last_attn_slice_mask is not None:\n",
    "                    new_attn_slice = torch.index_select(self.last_attn_slice, -1, self.last_attn_slice_indices)\n",
    "                    attn_slice = attn_slice * (1 - self.last_attn_slice_mask) + new_attn_slice * self.last_attn_slice_mask\n",
    "                else:\n",
    "                    attn_slice = self.last_attn_slice\n",
    "                \n",
    "                self.use_last_attn_slice = False\n",
    "                    \n",
    "            if self.save_last_attn_slice:\n",
    "                self.last_attn_slice = attn_slice\n",
    "                self.save_last_attn_slice = False\n",
    "                \n",
    "            if self.use_last_attn_weights and self.last_attn_slice_weights is not None:\n",
    "                attn_slice = attn_slice * self.last_attn_slice_weights\n",
    "                self.use_last_attn_weights = False\n",
    "            \n",
    "            attn_slice = torch.matmul(attn_slice, value[start_idx:end_idx])\n",
    "\n",
    "            hidden_states[start_idx:end_idx] = attn_slice\n",
    "\n",
    "        # reshape hidden_states\n",
    "        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)\n",
    "        return hidden_states\n",
    "\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\":\n",
    "            module.last_attn_slice = None\n",
    "            module.use_last_attn_slice = False\n",
    "            module.use_last_attn_weights = False\n",
    "            module.save_last_attn_slice = False\n",
    "            module._sliced_attention = new_sliced_attention.__get__(module, type(module))\n",
    "            module._attention = new_attention.__get__(module, type(module))\n",
    "            \n",
    "def use_last_tokens_attention(use=True):\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn2\" in name:\n",
    "            module.use_last_attn_slice = use\n",
    "            \n",
    "def use_last_tokens_attention_weights(use=True):\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn2\" in name:\n",
    "            module.use_last_attn_weights = use\n",
    "            \n",
    "def use_last_self_attention(use=True):\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn1\" in name:\n",
    "            module.use_last_attn_slice = use\n",
    "            \n",
    "def save_last_tokens_attention(save=True):\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn2\" in name:\n",
    "            module.save_last_attn_slice = save\n",
    "            \n",
    "def save_last_self_attention(save=True):\n",
    "    for name, module in unet.named_modules():\n",
    "        module_name = type(module).__name__\n",
    "        if module_name == \"CrossAttention\" and \"attn1\" in name:\n",
    "            module.save_last_attn_slice = save\n",
    "            \n",
    "@torch.no_grad()\n",
    "def stablediffusion(prompt=\"\", prompt_edit=None, prompt_edit_token_weights=[], prompt_edit_tokens_start=0.0, prompt_edit_tokens_end=1.0, prompt_edit_spatial_start=0.0, prompt_edit_spatial_end=1.0, guidance_scale=7.5, steps=50, seed=None, width=512, height=512, init_image=None, init_image_strength=0.5):\n",
    "    #Change size to multiple of 64 to prevent size mismatches inside model\n",
    "    width = width - width % 64\n",
    "    height = height - height % 64\n",
    "    \n",
    "    #If seed is None, randomly select seed from 0 to 2^32-1\n",
    "    if seed is None: seed = random.randrange(2**32 - 1)\n",
    "    generator = torch.cuda.manual_seed(seed)\n",
    "    \n",
    "    #Set inference timesteps to scheduler\n",
    "    scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n",
    "    scheduler.set_timesteps(steps)\n",
    "    \n",
    "    #Preprocess image if it exists (img2img)\n",
    "    if init_image is not None:\n",
    "        #Resize and transpose for numpy b h w c -> torch b c h w\n",
    "        init_image = init_image.resize((width, height), resample=Image.Resampling.LANCZOS)\n",
    "        init_image = np.array(init_image).astype(np.float32) / 255.0 * 2.0 - 1.0\n",
    "        init_image = torch.from_numpy(init_image[np.newaxis, ...].transpose(0, 3, 1, 2))\n",
    "        \n",
    "        #If there is alpha channel, composite alpha for white, as the diffusion model does not support alpha channel\n",
    "        if init_image.shape[1] > 3:\n",
    "            init_image = init_image[:, :3] * init_image[:, 3:] + (1 - init_image[:, 3:])\n",
    "            \n",
    "        #Move image to GPU\n",
    "        init_image = init_image.to(device)\n",
    "        \n",
    "        #Encode image\n",
    "        with autocast(device):\n",
    "            init_latent = vae.encode(init_image).latent_dist.sample(generator=generator) * 0.18215\n",
    "            \n",
    "        t_start = steps - int(steps * init_image_strength)\n",
    "            \n",
    "    else:\n",
    "        init_latent = torch.zeros((1, unet.in_channels, height // 8, width // 8), device=device)\n",
    "        t_start = 0\n",
    "    \n",
    "    #Generate random normal noise\n",
    "    noise = torch.randn(init_latent.shape, generator=generator, device=device)\n",
    "    #latent = noise * scheduler.init_noise_sigma\n",
    "    latent = scheduler.add_noise(init_latent, noise, torch.tensor([scheduler.timesteps[t_start]], device=device)).to(device)\n",
    "    \n",
    "    #Process clip\n",
    "    with autocast(device):\n",
    "        tokens_unconditional = clip_tokenizer(\"\", padding=\"max_length\", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors=\"pt\", return_overflowing_tokens=True)\n",
    "        embedding_unconditional = clip(tokens_unconditional.input_ids.to(device)).last_hidden_state\n",
    "\n",
    "        tokens_conditional = clip_tokenizer(prompt, padding=\"max_length\", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors=\"pt\", return_overflowing_tokens=True)\n",
    "        embedding_conditional = clip(tokens_conditional.input_ids.to(device)).last_hidden_state\n",
    "\n",
    "        #Process prompt editing\n",
    "        if prompt_edit is not None:\n",
    "            tokens_conditional_edit = clip_tokenizer(prompt_edit, padding=\"max_length\", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors=\"pt\", return_overflowing_tokens=True)\n",
    "            embedding_conditional_edit = clip(tokens_conditional_edit.input_ids.to(device)).last_hidden_state\n",
    "            \n",
    "            init_attention_edit(tokens_conditional, tokens_conditional_edit)\n",
    "            \n",
    "        init_attention_func()\n",
    "        init_attention_weights(prompt_edit_token_weights)\n",
    "            \n",
    "        timesteps = scheduler.timesteps[t_start:]\n",
    "        \n",
    "        for i, t in tqdm(enumerate(timesteps), total=len(timesteps)):\n",
    "            t_index = t_start + i\n",
    "\n",
    "            #sigma = scheduler.sigmas[t_index]\n",
    "            latent_model_input = latent\n",
    "            latent_model_input = scheduler.scale_model_input(latent_model_input, t)\n",
    "\n",
    "            #Predict the unconditional noise residual\n",
    "            noise_pred_uncond = unet(latent_model_input, t, encoder_hidden_states=embedding_unconditional).sample\n",
    "            \n",
    "            #Prepare the Cross-Attention layers\n",
    "            if prompt_edit is not None:\n",
    "                save_last_tokens_attention()\n",
    "                save_last_self_attention()\n",
    "            else:\n",
    "                #Use weights on non-edited prompt when edit is None\n",
    "                use_last_tokens_attention_weights()\n",
    "                \n",
    "            #Predict the conditional noise residual and save the cross-attention layer activations\n",
    "            noise_pred_cond = unet(latent_model_input, t, encoder_hidden_states=embedding_conditional).sample\n",
    "            \n",
    "            #Edit the Cross-Attention layer activations\n",
    "            if prompt_edit is not None:\n",
    "                t_scale = t / scheduler.num_train_timesteps\n",
    "                if t_scale >= prompt_edit_tokens_start and t_scale <= prompt_edit_tokens_end:\n",
    "                    use_last_tokens_attention()\n",
    "                if t_scale >= prompt_edit_spatial_start and t_scale <= prompt_edit_spatial_end:\n",
    "                    use_last_self_attention()\n",
    "                    \n",
    "                #Use weights on edited prompt\n",
    "                use_last_tokens_attention_weights()\n",
    "\n",
    "                #Predict the edited conditional noise residual using the cross-attention masks\n",
    "                noise_pred_cond = unet(latent_model_input, t, encoder_hidden_states=embedding_conditional_edit).sample\n",
    "                \n",
    "            #Perform guidance\n",
    "            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)\n",
    "\n",
    "            latent = scheduler.step(noise_pred, t_index, latent).prev_sample\n",
    "\n",
    "        #scale and decode the image latents with vae\n",
    "        latent = latent / 0.18215\n",
    "        image = vae.decode(latent.to(vae.dtype)).sample\n",
    "\n",
    "    image = (image / 2 + 0.5).clamp(0, 1)\n",
    "    image = image.cpu().permute(0, 2, 3, 1).numpy()\n",
    "    image = (image[0] * 255).round().astype(\"uint8\")\n",
    "    return Image.fromarray(image)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e9d39a-df21-45ba-bcfc-58b8784617d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prompt_token(prompt, index):\n",
    "    tokens = clip_tokenizer(prompt, padding=\"max_length\", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors=\"pt\", return_overflowing_tokens=True).input_ids[0]\n",
    "    return clip_tokenizer.decode(tokens[index:index+1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44ff9cb-6d43-4b32-bf2a-17bfca665218",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_token(\"a cat sitting on a car\", 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f69d449-0192-4db2-8687-9f9adab3907d",
   "metadata": {},
   "outputs": [],
   "source": [
    "stablediffusion(\"a cat sitting on a car\", seed=248396402679, steps=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba17dce4-f14c-4df3-8325-7adcabb33c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "stablediffusion(\"a cat sitting on a car\", \"a smiling dog sitting on a car\", prompt_edit_spatial_start=0.7, seed=248396402679)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c10c7e-38d4-4627-816d-241073684825",
   "metadata": {},
   "outputs": [],
   "source": [
    "stablediffusion(\"a cat sitting on a car\", \"a hamster sitting on a car\", prompt_edit_spatial_start=0.5, seed=248396402679)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b62f03c4-ee26-4f1d-a3c2-c83666e0a37e",
   "metadata": {},
   "outputs": [],
   "source": [
    "stablediffusion(\"a cat sitting on a car\", \"a tiger sitting on a car\", prompt_edit_spatial_start=1.0, seed=248396402679)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a16be9a-0411-4fbd-83f9-e968243d94fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54f31e7f-5527-4310-9f86-909c2f3525dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Remove fantasy\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", prompt_edit_token_weights=[(2, -8)], seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f919ec34-5038-4bcc-9bcf-c959b6d2e9f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Winter\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A winter fantasy landscape with a pine forest, trending on artstation\", seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e8bfb6c-4510-46e2-adfe-e99e38abf7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Watercolor style\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A watercolor painting of a landscape with a pine forest, trending on artstation\", seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995ea7e7-21c8-42ab-99db-3388469fe443",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Remove fog\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A fantasy landscape with a pine forest with fog, trending on artstation\", prompt_edit_token_weights=[(9, -6)], seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cbdf169-a0c3-47aa-8cf4-dee1312edab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Removing fog and rocks\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A fantasy landscape with a pine forest with fog and rocks, trending on artstation\", prompt_edit_token_weights=[(9, -6), (11, -6)], seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34d8cd67-a1ab-4572-82b1-337637b5cfce",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adding a river\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A fantasy landscape with a pine forest and a river, trending on artstation\", prompt_edit_spatial_end=0.8, seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57cf95b7-38f6-49f6-8492-266af0569e3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adding a lone cabin\n",
    "stablediffusion(\"A fantasy landscape with a pine forest, trending on artstation\", \"A fantasy landscape with a pine forest, lone cabin, trending on artstation\", prompt_edit_spatial_start=1.0, seed=2483964025, width=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb871a4d-153f-4afe-b6fe-a4cdd4635021",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "diffusers",
   "language": "python",
   "name": "diffusers"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}