Upload 19 files

.gitignore

@@ -0,0 +1,98 @@
+# Byte-compile / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.aider*
+
+# ignore ALL .log files
+*.log
+# ignore ALL files in ANY directory named temp
+.ipynb_checkpoints/
+images/.ipynb_checkpoints/
+results/
+.vscode/settings.json
+fonts/agency.ttf
+fonts/calibri.ttf
+data/

Comic_Generation.ipynb

@@ -0,0 +1,665 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "## StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation  \n",
+    "[![Paper page](https://huggingface.co/datasets/huggingface/badges/resolve/main/paper-page-md-dark.svg)]()\n",
+    "[[Paper]()] &emsp; [[Project Page]()] &emsp; <br>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Import Packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/tjut_lixiang/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "# %load_ext autoreload\n",
+    "# %autoreload 2\n",
+    "import gradio as gr\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import requests\n",
+    "import random\n",
+    "import os\n",
+    "import sys\n",
+    "import pickle\n",
+    "from PIL import Imagex\n",
+    "from tqdm.auto import tqdm\n",
+    "from datetime import datetime\n",
+    "from utils.gradio_utils import is_torch2_available\n",
+    "if is_torch2_available():\n",
+    "    from utils.gradio_utils import \\\n",
+    "        AttnProcessor2_0 as AttnProcessor\n",
+    "else:\n",
+    "    from utils.gradio_utils  import AttnProcessor\n",
+    "\n",
+    "import diffusers\n",
+    "from diffusers import StableDiffusionXLPipeline\n",
+    "from diffusers import DDIMScheduler\n",
+    "import torch.nn.functional as F\n",
+    "from utils.gradio_utils import cal_attn_mask_xl\n",
+    "import copy\n",
+    "import os\n",
+    "from diffusers.utils import load_image\n",
+    "from utils.utils import get_comic\n",
+    "from utils.style_template import styles"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Set Config "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Global\n",
+    "STYLE_NAMES = list(styles.keys())\n",
+    "DEFAULT_STYLE_NAME = \"(No style)\"\n",
+    "MAX_SEED = np.iinfo(np.int32).max\n",
+    "global models_dict\n",
+    "use_va = False\n",
+    "models_dict = {\n",
+    "   \"Juggernaut\":\"RunDiffusion/Juggernaut-XL-v8\",\n",
+    "   \"RealVision\":\"SG161222/RealVisXL_V4.0\" ,\n",
+    "   \"SDXL\":\"stabilityai/stable-diffusion-xl-base-1.0\" ,\n",
+    "   \"Unstable\": \"stablediffusionapi/sdxl-unstable-diffusers-y\"\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "torch.cuda.is_available()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def setup_seed(seed):\n",
+    "    torch.manual_seed(seed)\n",
+    "    torch.cuda.manual_seed_all(seed)\n",
+    "    np.random.seed(seed)\n",
+    "    random.seed(seed)\n",
+    "    torch.backends.cudnn.deterministic = True\n",
+    "\n",
+    "    \n",
+    "#################################################\n",
+    "########Consistent Self-Attention################\n",
+    "#################################################\n",
+    "class SpatialAttnProcessor2_0(torch.nn.Module):\n",
+    "    r\"\"\"\n",
+    "    Attention processor for IP-Adapater for PyTorch 2.0.\n",
+    "    Args:\n",
+    "        hidden_size (`int`):\n",
+    "            The hidden size of the attention layer.\n",
+    "        cross_attention_dim (`int`):\n",
+    "            The number of channels in the `encoder_hidden_states`.\n",
+    "        text_context_len (`int`, defaults to 77):\n",
+    "            The context length of the text features.\n",
+    "        scale (`float`, defaults to 1.0):\n",
+    "            the weight scale of image prompt.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self, hidden_size = None, cross_attention_dim=None,id_length = 4,device = \"cuda:0\",dtype = torch.float16):\n",
+    "        super().__init__()\n",
+    "        if not hasattr(F, \"scaled_dot_product_attention\"):\n",
+    "            raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n",
+    "        self.device = device\n",
+    "        self.dtype = dtype\n",
+    "        self.hidden_size = hidden_size\n",
+    "        self.cross_attention_dim = cross_attention_dim\n",
+    "        self.total_length = id_length + 1\n",
+    "        self.id_length = id_length\n",
+    "        self.id_bank = {}\n",
+    "\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        attn,\n",
+    "        hidden_states,\n",
+    "        encoder_hidden_states=None,\n",
+    "        attention_mask=None,\n",
+    "        temb=None):\n",
+    "        global total_count,attn_count,cur_step,mask1024,mask4096\n",
+    "        global sa32, sa64\n",
+    "        global write\n",
+    "        global height,width\n",
+    "        if write:\n",
+    "            # print(f\"white:{cur_step}\")\n",
+    "            self.id_bank[cur_step] = [hidden_states[:self.id_length], hidden_states[self.id_length:]]\n",
+    "        else:\n",
+    "            encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),hidden_states[:1],self.id_bank[cur_step][1].to(self.device),hidden_states[1:]))\n",
+    "        # skip in early step\n",
+    "        if cur_step <5:\n",
+    "            hidden_states = self.__call2__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)\n",
+    "        else:   # 256 1024 4096\n",
+    "            random_number = random.random()\n",
+    "            if cur_step <20:\n",
+    "                rand_num = 0.3\n",
+    "            else:\n",
+    "                rand_num = 0.1\n",
+    "            if random_number > rand_num:\n",
+    "                if not write:\n",
+    "                    if hidden_states.shape[1] == (height//32) * (width//32):\n",
+    "                        attention_mask = mask1024[mask1024.shape[0] // self.total_length * self.id_length:]\n",
+    "                    else:\n",
+    "                        attention_mask = mask4096[mask4096.shape[0] // self.total_length * self.id_length:]\n",
+    "                else:\n",
+    "                    if hidden_states.shape[1] == (height//32) * (width//32):\n",
+    "                        attention_mask = mask1024[:mask1024.shape[0] // self.total_length * self.id_length,:mask1024.shape[0] // self.total_length * self.id_length]\n",
+    "                    else:\n",
+    "                        attention_mask = mask4096[:mask4096.shape[0] // self.total_length * self.id_length,:mask4096.shape[0] // self.total_length * self.id_length]\n",
+    "                hidden_states = self.__call1__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)\n",
+    "            else:\n",
+    "                hidden_states = self.__call2__(attn, hidden_states,None,attention_mask,temb)\n",
+    "        attn_count +=1\n",
+    "        if attn_count == total_count:\n",
+    "            attn_count = 0\n",
+    "            cur_step += 1\n",
+    "            mask1024,mask4096 = cal_attn_mask_xl(self.total_length,self.id_length,sa32,sa64,height,width, device=self.device, dtype= self.dtype)\n",
+    "\n",
+    "        return hidden_states\n",
+    "    def __call1__(\n",
+    "        self,\n",
+    "        attn,\n",
+    "        hidden_states,\n",
+    "        encoder_hidden_states=None,\n",
+    "        attention_mask=None,\n",
+    "        temb=None,\n",
+    "    ):\n",
+    "        residual = hidden_states\n",
+    "        if attn.spatial_norm is not None:\n",
+    "            hidden_states = attn.spatial_norm(hidden_states, temb)\n",
+    "        input_ndim = hidden_states.ndim\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            total_batch_size, channel, height, width = hidden_states.shape\n",
+    "            hidden_states = hidden_states.view(total_batch_size, channel, height * width).transpose(1, 2)\n",
+    "        total_batch_size,nums_token,channel = hidden_states.shape\n",
+    "        img_nums = total_batch_size//2\n",
+    "        hidden_states = hidden_states.view(-1,img_nums,nums_token,channel).reshape(-1,img_nums * nums_token,channel)\n",
+    "\n",
+    "        batch_size, sequence_length, _ = hidden_states.shape\n",
+    "\n",
+    "        if attn.group_norm is not None:\n",
+    "            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n",
+    "\n",
+    "        query = attn.to_q(hidden_states)\n",
+    "\n",
+    "        if encoder_hidden_states is None:\n",
+    "            encoder_hidden_states = hidden_states  # B, N, C\n",
+    "        else:\n",
+    "            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,nums_token,channel).reshape(-1,(self.id_length+1) * nums_token,channel)\n",
+    "\n",
+    "        key = attn.to_k(encoder_hidden_states)\n",
+    "        value = attn.to_v(encoder_hidden_states)\n",
+    "\n",
+    "\n",
+    "        inner_dim = key.shape[-1]\n",
+    "        head_dim = inner_dim // attn.heads\n",
+    "\n",
+    "        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "\n",
+    "        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "        hidden_states = F.scaled_dot_product_attention(\n",
+    "            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n",
+    "        )\n",
+    "\n",
+    "        hidden_states = hidden_states.transpose(1, 2).reshape(total_batch_size, -1, attn.heads * head_dim)\n",
+    "        hidden_states = hidden_states.to(query.dtype)\n",
+    "\n",
+    "\n",
+    "\n",
+    "        # linear proj\n",
+    "        hidden_states = attn.to_out[0](hidden_states)\n",
+    "        # dropout\n",
+    "        hidden_states = attn.to_out[1](hidden_states)\n",
+    "\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            hidden_states = hidden_states.transpose(-1, -2).reshape(total_batch_size, channel, height, width)\n",
+    "        if attn.residual_connection:\n",
+    "            hidden_states = hidden_states + residual\n",
+    "        hidden_states = hidden_states / attn.rescale_output_factor\n",
+    "        # print(hidden_states.shape)\n",
+    "        return hidden_states\n",
+    "    def __call2__(\n",
+    "        self,\n",
+    "        attn,\n",
+    "        hidden_states,\n",
+    "        encoder_hidden_states=None,\n",
+    "        attention_mask=None,\n",
+    "        temb=None):\n",
+    "        residual = hidden_states\n",
+    "\n",
+    "        if attn.spatial_norm is not None:\n",
+    "            hidden_states = attn.spatial_norm(hidden_states, temb)\n",
+    "\n",
+    "        input_ndim = hidden_states.ndim\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            batch_size, channel, height, width = hidden_states.shape\n",
+    "            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n",
+    "\n",
+    "        batch_size, sequence_length, channel = (\n",
+    "            hidden_states.shape\n",
+    "        )\n",
+    "        # print(hidden_states.shape)\n",
+    "        if attention_mask is not None:\n",
+    "            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n",
+    "            # scaled_dot_product_attention expects attention_mask shape to be\n",
+    "            # (batch, heads, source_length, target_length)\n",
+    "            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n",
+    "\n",
+    "        if attn.group_norm is not None:\n",
+    "            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n",
+    "\n",
+    "        query = attn.to_q(hidden_states)\n",
+    "\n",
+    "        if encoder_hidden_states is None:\n",
+    "            encoder_hidden_states = hidden_states  # B, N, C\n",
+    "        else:\n",
+    "            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)\n",
+    "\n",
+    "        key = attn.to_k(encoder_hidden_states)\n",
+    "        value = attn.to_v(encoder_hidden_states)\n",
+    "\n",
+    "        inner_dim = key.shape[-1]\n",
+    "        head_dim = inner_dim // attn.heads\n",
+    "\n",
+    "        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "\n",
+    "        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "\n",
+    "        hidden_states = F.scaled_dot_product_attention(\n",
+    "            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n",
+    "        )\n",
+    "\n",
+    "        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n",
+    "        hidden_states = hidden_states.to(query.dtype)\n",
+    "\n",
+    "        # linear proj\n",
+    "        hidden_states = attn.to_out[0](hidden_states)\n",
+    "        # dropout\n",
+    "        hidden_states = attn.to_out[1](hidden_states)\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n",
+    "\n",
+    "        if attn.residual_connection:\n",
+    "            hidden_states = hidden_states + residual\n",
+    "\n",
+    "        hidden_states = hidden_states / attn.rescale_output_factor\n",
+    "\n",
+    "        return hidden_states\n",
+    "\n",
+    "def set_attention_processor(unet,id_length):\n",
+    "    attn_procs = {}\n",
+    "    for name in unet.attn_processors.keys():\n",
+    "        cross_attention_dim = None if name.endswith(\"attn1.processor\") else unet.config.cross_attention_dim\n",
+    "        if name.startswith(\"mid_block\"):\n",
+    "            hidden_size = unet.config.block_out_channels[-1]\n",
+    "        elif name.startswith(\"up_blocks\"):\n",
+    "            block_id = int(name[len(\"up_blocks.\")])\n",
+    "            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]\n",
+    "        elif name.startswith(\"down_blocks\"):\n",
+    "            block_id = int(name[len(\"down_blocks.\")])\n",
+    "            hidden_size = unet.config.block_out_channels[block_id]\n",
+    "        if cross_attention_dim is None:\n",
+    "            if name.startswith(\"up_blocks\") :\n",
+    "                attn_procs[name] = SpatialAttnProcessor2_0(id_length = id_length)\n",
+    "            else:    \n",
+    "                attn_procs[name] = AttnProcessor()\n",
+    "        else:\n",
+    "            attn_procs[name] = AttnProcessor()\n",
+    "\n",
+    "    unet.set_attn_processor(attn_procs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load Pipeline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/tjut_lixiang/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/huggingface_hub/file_download.py:1150: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
+      "  warnings.warn(\n",
+      "Loading pipeline components...: 100%|██████████| 7/7 [00:49<00:00,  7.13s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "successsfully load consistent self-attention\n",
+      "number of the processor : 36\n"
+     ]
+    }
+   ],
+   "source": [
+    "global attn_count, total_count, id_length, total_length,cur_step, cur_model_type\n",
+    "global write\n",
+    "global  sa32, sa64\n",
+    "global height,width\n",
+    "attn_count = 0\n",
+    "total_count = 0\n",
+    "cur_step = 0\n",
+    "id_length = 4\n",
+    "total_length = 5\n",
+    "cur_model_type = \"\"\n",
+    "device=\"cuda:0\"\n",
+    "global attn_procs,unet\n",
+    "attn_procs = {}\n",
+    "###\n",
+    "write = False\n",
+    "### strength of consistent self-attention: the larger, the stronger\n",
+    "sa32 = 0.5\n",
+    "sa64 = 0.5\n",
+    "### Res. of the Generated Comics. Please Note: SDXL models may do worse in a low-resolution! \n",
+    "height = 768\n",
+    "width = 768\n",
+    "###\n",
+    "global pipe\n",
+    "global sd_model_path\n",
+    "sd_model_path = models_dict[\"RealVision\"] #\"SG161222/RealVisXL_V4.0\"\n",
+    "### LOAD Stable Diffusion Pipeline\n",
+    "pipe = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16, use_safetensors=False)\n",
+    "pipe = pipe.to(device)\n",
+    "pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)\n",
+    "pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)\n",
+    "pipe.scheduler.set_timesteps(50)\n",
+    "unet = pipe.unet\n",
+    "\n",
+    "### Insert PairedAttention\n",
+    "for name in unet.attn_processors.keys():\n",
+    "    cross_attention_dim = None if name.endswith(\"attn1.processor\") else unet.config.cross_attention_dim\n",
+    "    if name.startswith(\"mid_block\"):\n",
+    "        hidden_size = unet.config.block_out_channels[-1]\n",
+    "    elif name.startswith(\"up_blocks\"):\n",
+    "        block_id = int(name[len(\"up_blocks.\")])\n",
+    "        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]\n",
+    "    elif name.startswith(\"down_blocks\"):\n",
+    "        block_id = int(name[len(\"down_blocks.\")])\n",
+    "        hidden_size = unet.config.block_out_channels[block_id]\n",
+    "    if cross_attention_dim is None and (name.startswith(\"up_blocks\") ) :\n",
+    "        attn_procs[name] =  SpatialAttnProcessor2_0(id_length = id_length)\n",
+    "        total_count +=1\n",
+    "    else:\n",
+    "        attn_procs[name] = AttnProcessor()\n",
+    "print(\"successsfully load consistent self-attention\")\n",
+    "print(f\"number of the processor : {total_count}\")\n",
+    "unet.set_attn_processor(copy.deepcopy(attn_procs))\n",
+    "global mask1024,mask4096\n",
+    "mask1024, mask4096 = cal_attn_mask_xl(total_length,id_length,sa32,sa64,height,width,device=device,dtype= torch.float16)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create the text description for the comics\n",
+    "Tips: Existing text2image diffusion models may not always generate images that accurately match text descriptions. Our training-free approach can improve the consistency of characters, but it does not enhance the control over the text. Therefore, in some cases, you may need to carefully craft your prompts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "guidance_scale = 5.0\n",
+    "seed = 2047\n",
+    "sa32 = 0.5\n",
+    "sa64 = 0.5\n",
+    "id_length = 4\n",
+    "num_steps = 50\n",
+    "general_prompt = \"a man with a black suit\"\n",
+    "negative_prompt = \"naked, deformed, bad anatomy, disfigured, poorly drawn face, mutation, extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, disconnected limbs, blurry, watermarks, oversaturated, distorted hands, amputation\"\n",
+    "prompt_array = [\"wake up in the bed\",\n",
+    "                \"have breakfast\",\n",
+    "                \"is on the road, go to the company\",\n",
+    "                \"work in the company\",\n",
+    "                \"running in the playground\",\n",
+    "                \"reading book in the home\"\n",
+    "                ]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      " 30%|███       | 15/50 [00:12<00:30,  1.16it/s]\n"
+     ]
+    },
+    {
+     "ename": "OutOfMemoryError",
+     "evalue": "CUDA out of memory. Tried to allocate 3.16 GiB (GPU 0; 23.70 GiB total capacity; 17.71 GiB already allocated; 1.04 GiB free; 21.02 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[7], line 20\u001b[0m\n\u001b[1;32m     18\u001b[0m attn_count \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     19\u001b[0m id_prompts, negative_prompt \u001b[38;5;241m=\u001b[39m apply_style(style_name, id_prompts, negative_prompt)\n\u001b[0;32m---> 20\u001b[0m id_images \u001b[38;5;241m=\u001b[39m \u001b[43mpipe\u001b[49m\u001b[43m(\u001b[49m\u001b[43mid_prompts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_inference_steps\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnum_steps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguidance_scale\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mguidance_scale\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[43mheight\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mheight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwidth\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mwidth\u001b[49m\u001b[43m,\u001b[49m\u001b[43mnegative_prompt\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnegative_prompt\u001b[49m\u001b[43m,\u001b[49m\u001b[43mgenerator\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mgenerator\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mimages\n\u001b[1;32m     22\u001b[0m write \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m id_image \u001b[38;5;129;01min\u001b[39;00m id_images:\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py:1216\u001b[0m, in \u001b[0;36mStableDiffusionXLPipeline.__call__\u001b[0;34m(self, prompt, prompt_2, height, width, num_inference_steps, timesteps, denoising_end, guidance_scale, negative_prompt, negative_prompt_2, num_images_per_prompt, eta, generator, latents, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, ip_adapter_image, output_type, return_dict, cross_attention_kwargs, guidance_rescale, original_size, crops_coords_top_left, target_size, negative_original_size, negative_crops_coords_top_left, negative_target_size, clip_skip, callback_on_step_end, callback_on_step_end_tensor_inputs, **kwargs)\u001b[0m\n\u001b[1;32m   1214\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m ip_adapter_image \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1215\u001b[0m     added_cond_kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mimage_embeds\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m image_embeds\n\u001b[0;32m-> 1216\u001b[0m noise_pred \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munet\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1217\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlatent_model_input\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1218\u001b[0m \u001b[43m    \u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1219\u001b[0m \u001b[43m    \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprompt_embeds\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1220\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtimestep_cond\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimestep_cond\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1221\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1222\u001b[0m \u001b[43m    \u001b[49m\u001b[43madded_cond_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madded_cond_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1223\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m   1224\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   1226\u001b[0m \u001b[38;5;66;03m# perform guidance\u001b[39;00m\n\u001b[1;32m   1227\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdo_classifier_free_guidance:\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/models/unet_2d_condition.py:1177\u001b[0m, in \u001b[0;36mUNet2DConditionModel.forward\u001b[0;34m(self, sample, timestep, encoder_hidden_states, class_labels, timestep_cond, attention_mask, cross_attention_kwargs, added_cond_kwargs, down_block_additional_residuals, mid_block_additional_residual, down_intrablock_additional_residuals, encoder_attention_mask, return_dict)\u001b[0m\n\u001b[1;32m   1174\u001b[0m     upsample_size \u001b[38;5;241m=\u001b[39m down_block_res_samples[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m2\u001b[39m:]\n\u001b[1;32m   1176\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(upsample_block, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_cross_attention\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;129;01mand\u001b[39;00m upsample_block\u001b[38;5;241m.\u001b[39mhas_cross_attention:\n\u001b[0;32m-> 1177\u001b[0m     sample \u001b[38;5;241m=\u001b[39m \u001b[43mupsample_block\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1178\u001b[0m \u001b[43m        \u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msample\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1179\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtemb\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43memb\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1180\u001b[0m \u001b[43m        \u001b[49m\u001b[43mres_hidden_states_tuple\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mres_samples\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1181\u001b[0m \u001b[43m        \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1182\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1183\u001b[0m \u001b[43m        \u001b[49m\u001b[43mupsample_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mupsample_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1184\u001b[0m \u001b[43m        \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1185\u001b[0m \u001b[43m        \u001b[49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1186\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1187\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1188\u001b[0m     sample \u001b[38;5;241m=\u001b[39m upsample_block(\n\u001b[1;32m   1189\u001b[0m         hidden_states\u001b[38;5;241m=\u001b[39msample,\n\u001b[1;32m   1190\u001b[0m         temb\u001b[38;5;241m=\u001b[39memb,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1193\u001b[0m         scale\u001b[38;5;241m=\u001b[39mlora_scale,\n\u001b[1;32m   1194\u001b[0m     )\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/models/unet_2d_blocks.py:2354\u001b[0m, in \u001b[0;36mCrossAttnUpBlock2D.forward\u001b[0;34m(self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states, cross_attention_kwargs, upsample_size, attention_mask, encoder_attention_mask)\u001b[0m\n\u001b[1;32m   2352\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2353\u001b[0m         hidden_states \u001b[38;5;241m=\u001b[39m resnet(hidden_states, temb, scale\u001b[38;5;241m=\u001b[39mlora_scale)\n\u001b[0;32m-> 2354\u001b[0m         hidden_states \u001b[38;5;241m=\u001b[39m \u001b[43mattn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2355\u001b[0m \u001b[43m            \u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2356\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2357\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2358\u001b[0m \u001b[43m            \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2359\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2360\u001b[0m \u001b[43m            \u001b[49m\u001b[43mreturn_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m   2361\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   2363\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupsamplers \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   2364\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m upsampler \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupsamplers:\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/models/transformer_2d.py:392\u001b[0m, in \u001b[0;36mTransformer2DModel.forward\u001b[0;34m(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, class_labels, cross_attention_kwargs, attention_mask, encoder_attention_mask, return_dict)\u001b[0m\n\u001b[1;32m    380\u001b[0m         hidden_states \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mutils\u001b[38;5;241m.\u001b[39mcheckpoint\u001b[38;5;241m.\u001b[39mcheckpoint(\n\u001b[1;32m    381\u001b[0m             create_custom_forward(block),\n\u001b[1;32m    382\u001b[0m             hidden_states,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    389\u001b[0m             \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mckpt_kwargs,\n\u001b[1;32m    390\u001b[0m         )\n\u001b[1;32m    391\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 392\u001b[0m         hidden_states \u001b[38;5;241m=\u001b[39m \u001b[43mblock\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    393\u001b[0m \u001b[43m            \u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    394\u001b[0m \u001b[43m            \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    395\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    396\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_attention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    397\u001b[0m \u001b[43m            \u001b[49m\u001b[43mtimestep\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimestep\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    398\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    399\u001b[0m \u001b[43m            \u001b[49m\u001b[43mclass_labels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclass_labels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    400\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    402\u001b[0m \u001b[38;5;66;03m# 3. Output\u001b[39;00m\n\u001b[1;32m    403\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_input_continuous:\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/models/attention.py:329\u001b[0m, in \u001b[0;36mBasicTransformerBlock.forward\u001b[0;34m(self, hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, timestep, cross_attention_kwargs, class_labels, added_cond_kwargs)\u001b[0m\n\u001b[1;32m    326\u001b[0m cross_attention_kwargs \u001b[38;5;241m=\u001b[39m cross_attention_kwargs\u001b[38;5;241m.\u001b[39mcopy() \u001b[38;5;28;01mif\u001b[39;00m cross_attention_kwargs \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m {}\n\u001b[1;32m    327\u001b[0m gligen_kwargs \u001b[38;5;241m=\u001b[39m cross_attention_kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgligen\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[0;32m--> 329\u001b[0m attn_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattn1\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    330\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnorm_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    331\u001b[0m \u001b[43m    \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43monly_cross_attention\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    332\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    333\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    334\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    335\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muse_ada_layer_norm_zero:\n\u001b[1;32m    336\u001b[0m     attn_output \u001b[38;5;241m=\u001b[39m gate_msa\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m1\u001b[39m) \u001b[38;5;241m*\u001b[39m attn_output\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/anaconda3/envs/storydiffusion/lib/python3.10/site-packages/diffusers/models/attention_processor.py:527\u001b[0m, in \u001b[0;36mAttention.forward\u001b[0;34m(self, hidden_states, encoder_hidden_states, attention_mask, **cross_attention_kwargs)\u001b[0m\n\u001b[1;32m    508\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;124;03mThe forward method of the `Attention` class.\u001b[39;00m\n\u001b[1;32m    510\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    522\u001b[0m \u001b[38;5;124;03m    `torch.Tensor`: The output of the attention layer.\u001b[39;00m\n\u001b[1;32m    523\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    524\u001b[0m \u001b[38;5;66;03m# The `Attention` class can call different attention processors / attention functions\u001b[39;00m\n\u001b[1;32m    525\u001b[0m \u001b[38;5;66;03m# here we simply pass along all tensors to the selected processor class\u001b[39;00m\n\u001b[1;32m    526\u001b[0m \u001b[38;5;66;03m# For standard processors that are defined here, `**cross_attention_kwargs` is empty\u001b[39;00m\n\u001b[0;32m--> 527\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocessor\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    528\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    529\u001b[0m \u001b[43m    \u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    530\u001b[0m \u001b[43m    \u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    531\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    532\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcross_attention_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    533\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[4], line 74\u001b[0m, in \u001b[0;36mSpatialAttnProcessor2_0.__call__\u001b[0;34m(self, attn, hidden_states, encoder_hidden_states, attention_mask, temb)\u001b[0m\n\u001b[1;32m     72\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     73\u001b[0m             attention_mask \u001b[38;5;241m=\u001b[39m mask4096[:mask4096\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtotal_length \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mid_length,:mask4096\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtotal_length \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mid_length]\n\u001b[0;32m---> 74\u001b[0m     hidden_states \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__call1__\u001b[49m\u001b[43m(\u001b[49m\u001b[43mattn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[43m,\u001b[49m\u001b[43mencoder_hidden_states\u001b[49m\u001b[43m,\u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtemb\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     76\u001b[0m     hidden_states \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__call2__(attn, hidden_states,\u001b[38;5;28;01mNone\u001b[39;00m,attention_mask,temb)\n",
+      "Cell \u001b[0;32mIn[4], line 127\u001b[0m, in \u001b[0;36mSpatialAttnProcessor2_0.__call1__\u001b[0;34m(self, attn, hidden_states, encoder_hidden_states, attention_mask, temb)\u001b[0m\n\u001b[1;32m    125\u001b[0m key \u001b[38;5;241m=\u001b[39m key\u001b[38;5;241m.\u001b[39mview(batch_size, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, attn\u001b[38;5;241m.\u001b[39mheads, head_dim)\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m    126\u001b[0m value \u001b[38;5;241m=\u001b[39m value\u001b[38;5;241m.\u001b[39mview(batch_size, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, attn\u001b[38;5;241m.\u001b[39mheads, head_dim)\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m)\n\u001b[0;32m--> 127\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscaled_dot_product_attention\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    128\u001b[0m \u001b[43m    \u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mattn_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdropout_p\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mis_causal\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\n\u001b[1;32m    129\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    131\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m hidden_states\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m)\u001b[38;5;241m.\u001b[39mreshape(total_batch_size, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, attn\u001b[38;5;241m.\u001b[39mheads \u001b[38;5;241m*\u001b[39m head_dim)\n\u001b[1;32m    132\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m hidden_states\u001b[38;5;241m.\u001b[39mto(query\u001b[38;5;241m.\u001b[39mdtype)\n",
+      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 3.16 GiB (GPU 0; 23.70 GiB total capacity; 17.71 GiB already allocated; 1.04 GiB free; 21.02 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
+     ]
+    }
+   ],
+   "source": [
+    "##########################################################################################\n",
+    "def apply_style_positive(style_name: str, positive: str):\n",
+    "    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])\n",
+    "    return p.replace(\"{prompt}\", positive) \n",
+    "def apply_style(style_name: str, positives: list, negative: str = \"\"):\n",
+    "    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])\n",
+    "    return [p.replace(\"{prompt}\", positive) for positive in positives], n + ' ' + negative\n",
+    "### Set the generated Style\n",
+    "style_name = \"Comic book\"\n",
+    "setup_seed(seed)\n",
+    "generator = torch.Generator(device=\"cuda:0\").manual_seed(seed)\n",
+    "prompts = [general_prompt+\",\"+prompt for prompt in prompt_array]\n",
+    "id_prompts = prompts[:id_length]\n",
+    "real_prompts = prompts[id_length:]\n",
+    "torch.cuda.empty_cache()\n",
+    "write = True\n",
+    "cur_step = 0\n",
+    "attn_count = 0\n",
+    "id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)\n",
+    "id_images = pipe(id_prompts, num_inference_steps = num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images\n",
+    "\n",
+    "write = False\n",
+    "for id_image in id_images:\n",
+    "    display(id_image)\n",
+    "real_images = []\n",
+    "for real_prompt in real_prompts:\n",
+    "    cur_step = 0\n",
+    "    real_prompt = apply_style_positive(style_name, real_prompt)\n",
+    "    real_images.append(pipe(real_prompt,  num_inference_steps=num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])\n",
+    "for real_image in real_images:\n",
+    "    display(real_image)  "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Continued Creation\n",
+    "From now on, you can create endless stories about this character without worrying about memory constraints."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "new_prompt_array = [\"siting on the sofa\",\n",
+    "             \"on the bed, at night \"]\n",
+    "new_prompts = [general_prompt+\",\"+prompt for prompt in new_prompt_array]\n",
+    "new_images = []\n",
+    "for new_prompt in new_prompts :\n",
+    "    cur_step = 0\n",
+    "    new_prompt = apply_style_positive(style_name, new_prompt)\n",
+    "    new_images.append(pipe(new_prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])\n",
+    "for new_image in new_images:\n",
+    "    display(new_image)  "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Make pictures into comics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 222,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "###\n",
+    "total_images = id_images + real_images + new_images\n",
+    "from PIL import Image,ImageOps,ImageDraw, ImageFont\n",
+    "#### LOAD Fonts, can also replace with any Fonts you have!\n",
+    "font = ImageFont.truetype(\"./fonts/Inkfree.ttf\", 30)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 223,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import importlib\n",
+    "# import utils.utils\n",
+    "# importlib.reload(utils)\n",
+    "from utils.utils import get_row_image\n",
+    "from utils.utils import get_row_image\n",
+    "from utils.utils import get_comic_4panel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "comics = get_comic_4panel(total_images, captions = prompt_array+ new_prompts,font = font )\n",
+    "for comic in comics:\n",
+    "    display(comic)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "fileId": "51613593-0d85-430e-8fce-c85e580fc483",
+  "kernelspec": {
+   "display_name": "storydiffusion",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,152 @@
+<p align="center">
+  <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/f79da6b7-0b3b-4dd7-8dd0-ba0b15306fe6" height=100>
+</p>
+
+<div align="center">
+  
+## StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation  [![Paper page](https://huggingface.co/datasets/huggingface/badges/resolve/main/paper-page-md-dark.svg)]()
+
+[[Paper](https://arxiv.org/abs/2405.01434)] &emsp; [[Project Page](https://storydiffusion.github.io/)] &emsp;  [[Jittor Version](https://github.com/JittorCV/jittordiffusion/tree/master)]&emsp; [[🤗 Comic Generation Demo ](https://huggingface.co/spaces/YupengZhou/StoryDiffusion)] [![Replicate](https://replicate.com/cjwbw/StoryDiffusion/badge)](https://replicate.com/cjwbw/StoryDiffusion) [![Run Comics Demo in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/HVision-NKU/StoryDiffusion/blob/main/Comic_Generation.ipynb) <br>
+</div>
+
+
 ---
-title: Demo
-emoji: 📉
-colorFrom: blue
-colorTo: red
-sdk: gradio
-sdk_version: 5.31.0
-app_file: app.py
-pinned: false
+
+Official implementation of **[StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation]()**.
+
+### **Demo Video**
+
+https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/d5b80f8f-09b0-48cd-8b10-daff46d422af
+
+
+### Update History
+
+***You can visit [here](update.md) to visit update history.***
+
+### 🌠  **Key Features:**
+StoryDiffusion can create a magic story by generating consistent images and videos. Our work mainly has two parts: 
+1. Consistent self-attention for character-consistent image generation over long-range sequences. It is hot-pluggable and compatible with all SD1.5 and SDXL-based image diffusion models. For the current implementation, the user needs to provide at least 3 text prompts for the consistent self-attention module. We recommend at least 5 - 6 text prompts for better layout arrangement.
+2. Motion predictor for long-range video generation, which predicts motion between Condition Images in a compressed image semantic space, achieving larger motion prediction. 
+
+
+
+## 🔥 **Examples**
+
+
+### Comics generation 
+
+
+![1](https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/b3771cbc-b6ca-4e26-bdc5-d944daf9f266)
+
+
+
+### Image-to-Video generation （Results are HIGHLY compressed for speed）
+Leveraging the images produced through our Consistent Self-Attention mechanism, we can extend the process to create videos by seamlessly transitioning between these images. This can be considered as a two-stage long video generation approach.
+
+Note: results are **highly compressed** for speed, you can visit [our website](https://storydiffusion.github.io/) for the high-quality version.
+#### Two-stage Long Videos Generation (New Update)
+Combining the two parts, we can generate very long and high-quality AIGC videos.
+| Video1 | Video2  | Video3  |
+| --- | --- | --- |
+| <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/4e7e0f24-5f90-419b-9a1e-cdf36d361b26" width=224>  | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/f509343d-d691-4e2a-b615-7d96381ef7c1" width=224> | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/4f0f7abb-4ae4-47a6-b692-5bdd8d9c8006" width=224>  |
+
+
+#### Long Video Results using Condition Images
+Our Image-to-Video model can generate a video by providing a sequence of user-input condition images.
+| Video1 | Video2  | Video3  |
+| --- | --- | --- |
+| <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/af6f5c50-c773-4ef2-a757-6d7a46393f39" width=224>  | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/d58e4037-d8df-4f90-8c81-ce4b6d2d868e" width=224> |  <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/40da15ba-f5c1-48d8-84d6-8d327207d696" width=224>  |
+
+| Video4 | Video5  | Video6  |
+| --- | --- | --- |
+| <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/8f04c9fc-3031-49e3-9de8-83d582b80a1f" width=224>  | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/604107fb-8afe-4052-bda4-362c646a756e" width=224> |  <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/b05fa6a0-12e6-4111-abf8-18b8cd84f3ff" width=224>  |
+
+
+
+
+#### Short Videos 
+
+| Video1 | Video2  | Video3  |
+| --- | --- | --- |
+| <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/5e7f717f-daad-46f6-b3ba-c087bd843158" width=224>  | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/79aa52b2-bf37-4c9c-8555-c7050aec0cdf" width=224> | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/9fdfd091-10e6-434e-9ce7-6d6e6d8f4b22" width=224>  |
+
+
+
+| Video4 | Video5  | Video6  |
+| --- | --- | --- |
+| <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/0b219b60-a998-4820-9657-6abe1747cb6b" width=224>  | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/d387aef0-ffc8-41b0-914f-4b0392d9f8c5" width=224> | <img src="https://github.com/HVision-NKU/StoryDiffusion/assets/49511209/3c64958a-1079-4ca0-a9cf-e0486adbc57f" width=224>  |
+
+
+
+
+## 🚩 **TODO/Updates**
+- [x] Comic Results of StoryDiffusion.
+- [x] Video Results of StoryDiffusion.
+- [x] Source code of Comic Generation
+- [x] Source code of gradio demo
+- [ ] Source code of Video Generation Model
+- [ ] Pretrained weight of Video Generation Model
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🔧 Dependencies and Installation
+
+- Python >= 3.8 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html))
+- [PyTorch >= 2.0.0](https://pytorch.org/)
+```bash
+conda create --name storydiffusion python=3.10
+conda activate storydiffusion
+pip install -U pip
+
+# Install requirements
+pip install -r requirements.txt
+```
+# How to use
+
+Currently, we provide two ways for you to generate comics.
+
+## Use the jupyter notebook
+
+You can open the `Comic_Generation.ipynb` and run the code.
+
+## Start a local gradio demo
+Run the following command:
+
+
+**(Recommend)** We provide a low GPU Memory cost version, it was tested on a machine with 24GB GPU-memory(Tesla A10) and 30GB RAM, and expected to work well with >20 G GPU-memory.
+
+```python
+python gradio_app_sdxl_specific_id_low_vram.py
+```
+
+
+## Contact
+If you have any questions, you are very welcome to email ypzhousdu@gmail.com and zhoudaquan21@gmail.com
+
+   
+
+
+# Disclaimer
+This project strives to impact the domain of AI-driven image and video generation positively. Users are granted the freedom to create images and videos using this tool, but they are expected to comply with local laws and utilize it responsibly. The developers do not assume any responsibility for potential misuse by users.
+
+# Related Resources
+Following are some third-party implementations of StoryDiffusion.
+
+
+## API
+
+- [runpod.io serverless worker](https://github.com/bes-dev/story-diffusion-runpod-serverless-worker) provided by [BeS](https://github.com/bes-dev).
+- [Replicate worker](https://github.com/camenduru/StoryDiffusion-replicate) provided by [camenduru](https://github.com/camenduru).
+
+
+
+
+# BibTeX
+If you find StoryDiffusion useful for your research and applications, please cite using this BibTeX:
+
+```BibTeX
+@article{zhou2024storydiffusion,
+  title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+  author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+  journal={NeurIPS 2024},
+  year={2024}
+}

app.py

@@ -0,0 +1,750 @@
+from email.policy import default
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import requests
+import random
+import os
+import sys
+import pickle
+from PIL import Image
+from tqdm.auto import tqdm
+from datetime import datetime
+from utils.gradio_utils import is_torch2_available
+if is_torch2_available():
+    from utils.gradio_utils import \
+        AttnProcessor2_0 as AttnProcessor
+    # from utils.gradio_utils import SpatialAttnProcessor2_0
+else:
+    from utils.gradio_utils  import AttnProcessor
+
+import diffusers
+from diffusers import StableDiffusionXLPipeline
+from utils import PhotoMakerStableDiffusionXLPipeline
+from diffusers import DDIMScheduler
+import torch.nn.functional as F
+from utils.gradio_utils import cal_attn_mask_xl
+import copy
+import os
+from huggingface_hub import hf_hub_download
+from diffusers.utils import load_image
+from utils.utils import get_comic
+from utils.style_template import styles
+image_encoder_path = "./data/models/ip_adapter/sdxl_models/image_encoder"
+ip_ckpt = "./data/models/ip_adapter/sdxl_models/ip-adapter_sdxl_vit-h.bin"
+os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Japanese Anime"
+global models_dict
+use_va = True
+models_dict = {
+#    "Juggernaut": "RunDiffusion/Juggernaut-XL-v8",
+#    "RealVision": "SG161222/RealVisXL_V4.0" ,
+#    "SDXL":"stabilityai/stable-diffusion-xl-base-1.0" ,
+   "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y"
+}
+photomaker_path =  hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")
+MAX_SEED = np.iinfo(np.int32).max
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+def set_text_unfinished():
+    return gr.update(visible=True, value="<h3>(Not Finished) Generating ···  The intermediate results will be shown.</h3>")
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>Generation Finished</h3>")
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])
+    return image_path_list
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+################################################################################################################################################################################
+    def __init__(self, hidden_size = None, cross_attention_dim=None,id_length = 4,device = "cuda",dtype = torch.float16):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count,attn_count,cur_step,mask1024,mask4096
+        global sa32, sa64
+        global write
+        global height,width
+        if write:
+            # print(f"white:{cur_step}")
+            self.id_bank[cur_step] = [hidden_states[:self.id_length], hidden_states[self.id_length:]]
+        else:
+            encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),hidden_states[:1],self.id_bank[cur_step][1].to(self.device),hidden_states[1:]))
+        # 判断随机数是否大于0.5
+        if cur_step <5:
+            hidden_states = self.__call2__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
+        else:   # 256 1024 4096
+            random_number = random.random()
+            if cur_step <20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                # print("mask shape",mask1024.shape,mask4096.shape)
+                if not write:
+                    if hidden_states.shape[1] == (height//32) * (width//32):
+                        attention_mask = mask1024[mask1024.shape[0] // self.total_length * self.id_length:]
+                    else:
+                        attention_mask = mask4096[mask4096.shape[0] // self.total_length * self.id_length:]
+                else:
+                    # print(self.total_length,self.id_length,hidden_states.shape,(height//32) * (width//32))
+                    if hidden_states.shape[1] == (height//32) * (width//32):
+                        attention_mask = mask1024[:mask1024.shape[0] // self.total_length * self.id_length,:mask1024.shape[0] // self.total_length * self.id_length]
+                    else:
+                        attention_mask = mask4096[:mask4096.shape[0] // self.total_length * self.id_length,:mask4096.shape[0] // self.total_length * self.id_length]
+                   # print(attention_mask.shape)
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                hidden_states = self.__call1__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
+            else:
+                hidden_states = self.__call2__(attn, hidden_states,None,attention_mask,temb)
+        attn_count +=1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            mask1024,mask4096 = cal_attn_mask_xl(self.total_length,self.id_length,sa32,sa64,height,width, device=self.device, dtype= self.dtype)
+
+        return hidden_states
+    def __call1__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # print("hidden state shape",hidden_states.shape,self.id_length)
+        residual = hidden_states
+        # if encoder_hidden_states is not None:
+        #     raise Exception("not implement")
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            total_batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(total_batch_size, channel, height * width).transpose(1, 2)
+        total_batch_size,nums_token,channel = hidden_states.shape
+        img_nums = total_batch_size//2
+        hidden_states = hidden_states.view(-1,img_nums,nums_token,channel).reshape(-1,img_nums * nums_token,channel)
+
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,nums_token,channel).reshape(-1,(self.id_length+1) * nums_token,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        # print(key.shape,value.shape,query.shape,attention_mask.shape)
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        #print(query.shape,key.shape,value.shape,attention_mask.shape)
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(total_batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        # if input_ndim == 4:
+        #     tile_hidden_states = tile_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        # if attn.residual_connection:
+        #     tile_hidden_states = tile_hidden_states + residual
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(total_batch_size, channel, height, width)
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        hidden_states = hidden_states / attn.rescale_output_factor
+        # print(hidden_states.shape)
+        return hidden_states
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, channel = (
+            hidden_states.shape
+        )
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+def set_attention_processor(unet,id_length,is_ipadapter = False):
+    global total_count
+    total_count = 0
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks") :
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length = id_length)
+                total_count +=1
+            else:    
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+    print("successsfully load paired self-attention")
+    print(f"number of the processor : {total_count}")
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = '''
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+'''
+
+
+#################################################
+title = r"""
+<h1 align="center">StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</h1>
+"""
+
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+3️⃣ Choose your preferred style template.<br>
+4️⃣ Click the <b>Submit</b> button to start customizing.
+"""
+
+article = r"""
+
+If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+---
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+
+```bibtex
+@article{Zhou2024storydiffusion,
+  title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+  author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+  year={2024}
+}
+```
+📋 **License**
+<br>
+The Contents you create are under Apache-2.0 LICENSE. The Code are under Attribution-NonCommercial 4.0 International. 
+
+📧 **Contact**
+<br>
+If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+"""
+version = r"""
+<h3 align="center">StoryDiffusion Version 0.01 (test version)</h3>
+
+<h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+<h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+<h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling."),Currently, support is only using Textual Description</h5>
+<h5 align="center">Tips: Not Ready Now! Just Test</h5>
+"""
+#################################################
+global attn_count, total_count, id_length, total_length,cur_step, cur_model_type
+global write
+global  sa32, sa64
+global height,width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+device="cuda"
+global attn_procs,unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global sd_model_path
+sd_model_path = models_dict["Unstable"]#"SG161222/RealVisXL_V4.0"
+use_safetensors= False
+### LOAD Stable Diffusion Pipeline
+pipe1 = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16, use_safetensors= use_safetensors)
+pipe1 = pipe1.to("cuda")
+pipe1.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe1.scheduler.set_timesteps(50)
+### 
+pipe2 = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+    sd_model_path, torch_dtype=torch.float16, use_safetensors=use_safetensors)
+pipe2 = pipe2.to("cuda")
+pipe2.load_photomaker_adapter(
+    os.path.dirname(photomaker_path),
+    subfolder="",
+    weight_name=os.path.basename(photomaker_path),
+    trigger_word="img"  # define the trigger word
+)
+pipe2 = pipe2.to("cuda")
+pipe2.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+pipe2.fuse_lora()
+
+######### Gradio Fuction #############
+
+def swap_to_gallery(images):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+def remove_tips():
+    return gr.update(visible=False)
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive) 
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [p.replace("{prompt}", positive) for positive in positives], n + ' ' + negative
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    elif _model_type == "Using Ref Images":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
+    else:
+        raise ValueError("Invalid model type",_model_type)
+
+
+######### Image Generation ##############
+@spaces.GPU
+def process_generation(_sd_type,_model_type,_upload_images, _num_steps,style_name, _Ip_Adapter_Strength ,_style_strength_ratio, guidance_scale, seed_,  sa32_, sa64_, id_length_,  general_prompt, negative_prompt,prompt_array,G_height,G_width,_comic_type):
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error("Please add the triger word \" img \"  behind the class word you want to customize, such as: man img or woman img")
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64,id_length,total_length,attn_procs,unet,cur_model_type,device
+    global write
+    global cur_step,attn_count
+    global height,width
+    height = G_height
+    width = G_width
+    global pipe1,pipe2
+    global sd_model_path,models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    if _model_type == "original":
+        pipe = pipe1
+        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    elif _model_type == "Photomaker":
+        pipe = pipe2
+        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    else:
+        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+        ##### ########################
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+    cur_model_type = _sd_type+"-"+_model_type+""+str(id_length_)
+    if _model_type != "original":
+        input_id_images = []
+        for img in _upload_images:
+            print(img)
+            input_id_images.append(load_image(img))
+    prompts = prompt_array.splitlines()
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device="cuda").manual_seed(seed_)
+    sa32, sa64 =  sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    prompts = [general_prompt + "," + prompt if "[NC]" not in prompt else prompt.replace("[NC]","")  for prompt in clipped_prompts]
+    prompts = [prompt.rpartition('#')[0] if "#" in prompt else prompt for prompt in prompts]
+    print(prompts)
+    id_prompts = prompts[:id_length]
+    real_prompts = prompts[id_length:]
+    torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    setup_seed(seed_)
+    total_results = []
+    if _model_type == "original":
+        id_images = pipe(id_prompts, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
+    elif _model_type == "Photomaker":
+        id_images = pipe(id_prompts,input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale, start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
+    else: 
+        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+    total_results = id_images + total_results
+    yield total_results
+    real_images = []
+    write = False
+    for real_prompt in real_prompts:
+        setup_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":   
+            real_images.append(pipe(real_prompt,  num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
+        elif _model_type == "Photomaker":      
+            real_images.append(pipe(real_prompt, input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
+        else:
+            raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+        total_results = [real_images[-1]] + total_results
+        yield total_results
+    if _comic_type != "No typesetting (default)":
+        captions= prompt_array.splitlines()
+        captions = [caption.replace("[NC]","") for caption in captions]
+        captions = [caption.split('#')[-1] if "#" in caption else caption for caption in captions]
+        from PIL import ImageFont
+        total_results = get_comic(id_images + real_images, _comic_type,captions= captions,font=ImageFont.truetype("./fonts/Inkfree.ttf", int(45))) + total_results
+    set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    yield total_results
+
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i,part in enumerate(arr):
+        if i != len(arr)-1:
+            stringtmp += part +"\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+with gr.Blocks(css=css) as demo:
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+
+    # gr.Markdown(logo)
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Group(elem_id="main-image"):
+            # button_run = gr.Button("generate id images ! 😺", elem_id="main_button", interactive=True)
+
+            prompts = []
+            colors = []
+            # with gr.Column(visible=False) as post_sketch:
+            #     for n in range(MAX_COLORS):
+            #         if n == 0 :
+            #             with gr.Row(visible=False) as color_row[n]:
+            #                 colors.append(gr.Image(shape=(100, 100), label="background", type="pil", image_mode="RGB", width=100, height=100))
+            #                 prompts.append(gr.Textbox(label="Prompt for the background (white region)", value=""))
+            #         else:
+            #             with gr.Row(visible=False) as color_row[n]:
+            #                 colors.append(gr.Image(shape=(100, 100), label="segment "+str(n), type="pil", image_mode="RGB", width=100, height=100))
+            #                 prompts.append(gr.Textbox(label="Prompt for the segment "+str(n)))
+
+            #     get_genprompt_run = gr.Button("(2) I've finished segment labeling ! 😺", elem_id="prompt_button", interactive=True)
+
+            with gr.Column(visible=True) as gen_prompt_vis:
+                sd_type = gr.Dropdown(choices=list(models_dict.keys()), value = "Unstable",label="sd_type", info="Select pretrained model")
+                model_type = gr.Radio(["Only Using Textual Description", "Using Ref Images"], label="model_type", value = "Only Using Textual Description",  info="Control type of the Character")
+                with gr.Group(visible=False) as control_image_input:
+                    files = gr.Files(
+                                label="Drag (Select) 1 or more photos of your face",
+                                file_types=["image"],
+                            )
+                    uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
+                    with gr.Column(visible=False) as clear_button:
+                        remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+                general_prompt = gr.Textbox(value='', label="(1) Textual Description for Character", interactive=True)
+                negative_prompt = gr.Textbox(value='', label="(2) Negative_prompt", interactive=True)
+                style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                prompt_array = gr.Textbox(lines = 3,value='', label="(3) Comic Description (each line corresponds to a frame).", interactive=True)
+                with gr.Accordion("(4) Tune the hyperparameters", open=True):
+                    #sa16_ = gr.Slider(label=" (The degree of Paired Attention at 16 x 16 self-attention layers) ", minimum=0, maximum=1., value=0.3, step=0.1)
+                    sa32_ = gr.Slider(label=" (The degree of Paired Attention at 32 x 32 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
+                    sa64_ = gr.Slider(label=" (The degree of Paired Attention at 64 x 64 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
+                    id_length_ = gr.Slider(label= "Number of id images in total images" , minimum=2, maximum=4, value=2, step=1)
+                    # total_length_ = gr.Slider(label= "Number of total images", minimum=1, maximum=20, value=1, step=1)
+                    seed_ = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, value=0, step=1)
+                    num_steps = gr.Slider( 
+                        label="Number of sample steps",
+                        minimum=20,
+                        maximum=100,
+                        step=1,
+                        value=50,
+                    )
+                    G_height = gr.Slider( 
+                        label="height",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    G_width = gr.Slider( 
+                        label="width",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    comic_type = gr.Radio(["No typesetting (default)", "Four Pannel", "Classic Comic Style"], value = "Classic Comic Style", label="Typesetting Style", info="Select the typesetting style ")
+                    guidance_scale = gr.Slider(
+                        label="Guidance scale",
+                        minimum=0.1,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5,
+                    )
+                    style_strength_ratio = gr.Slider(
+                        label="Style strength of Ref Image (%)",
+                        minimum=15,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                        visible=False
+                    )
+                    Ip_Adapter_Strength = gr.Slider(
+                        label="Ip_Adapter_Strength",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=0.5,
+                        visible=False
+                    )
+                final_run_btn = gr.Button("Generate ! 😺")
+
+
+        with gr.Column():
+            out_image = gr.Gallery(label="Result", columns=2, height='auto')
+            generated_information = gr.Markdown(label="Generation Details", value="",visible=False)
+            gr.Markdown(version)
+    model_type.change(fn = change_visiale_by_model_type , inputs = model_type, outputs=[control_image_input,style_strength_ratio,Ip_Adapter_Strength])
+    files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
+    remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
+
+    final_run_btn.click(fn=set_text_unfinished, outputs = generated_information
+    ).then(process_generation, inputs=[sd_type,model_type,files, num_steps,style, Ip_Adapter_Strength,style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array,G_height,G_width,comic_type], outputs=out_image
+    ).then(fn=set_text_finished,outputs = generated_information)
+
+
+    gr.Examples(
+        examples=[
+            [1,0.5,0.5,3,"a woman img, wearing a white T-shirt, blue loose hair",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                   array2string(["wake up in the bed",
+                                "have breakfast",
+                                "is on the road, go to company",
+                                "work in the company",
+                                "Take a walk next to the company at noon",
+                                "lying in bed at night"]),
+                                "Japanese Anime",  "Using Ref Images",get_image_path_list('./examples/taylor'),768,768
+                ],
+                [0,0.5,0.5,2,"a man, wearing black jacket",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                   array2string(["wake up in the bed",
+                                "have breakfast",
+                                "is on the road, go to the company,  close look",
+                                "work in the company",
+                                "laughing happily",
+                                "lying in bed at night"
+                                ]),
+                                "Japanese Anime","Only Using Textual Description",get_image_path_list('./examples/taylor'),768,768
+                ],
+                [0,0.3,0.5,2,"a girl, wearing white shirt, black skirt, black tie, yellow hair",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                    array2string([
+                            "at home #at home, began to go to drawing",
+                            "sitting alone on a park bench.",
+                            "reading a book on a park bench.",
+                            "[NC]A squirrel approaches, peeking over the bench. ",
+                            "look around in the park. # She looks around and enjoys the beauty of nature.",
+                            "[NC]leaf falls from the tree, landing on the sketchbook.",
+                            "picks up the leaf, examining its details closely.",
+                            "starts sketching the leaf with intricate lines.",
+                            "holds up the sketch drawing of the leaf.",
+                            "[NC]The brown squirrel appear.",
+                            "is very happy # She is very happy to see the squirrel again",
+                            "[NC]The brown squirrel takes the cracker and scampers up a tree. # She gives the squirrel cracker",
+                            "laughs and tucks the leaf into her book as a keepsake.",
+                            "ready to leave.",]),
+                    "Japanese Anime","Only Using Textual Description",get_image_path_list('./examples/taylor'),768,768
+                ]
+                ],
+        inputs=[seed_, sa32_, sa64_, id_length_,  general_prompt, negative_prompt, prompt_array,style,model_type,files,G_height,G_width],
+        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
+        # run_on_click=True,
+        label='😺 Examples 😺',
+    )
+    gr.Markdown(article)
+
+    # demo.load(None, None, None, _js=load_js)
+
+demo.launch(server_name="0.0.0.0", share = True if use_va else False)

app1.py

@@ -0,0 +1,750 @@
+from email.policy import default
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import requests
+import random
+import os
+import sys
+import pickle
+from PIL import Image
+from tqdm.auto import tqdm
+from datetime import datetime
+from utils.gradio_utils import is_torch2_available
+if is_torch2_available():
+    from utils.gradio_utils import \
+        AttnProcessor2_0 as AttnProcessor
+    # from utils.gradio_utils import SpatialAttnProcessor2_0
+else:
+    from utils.gradio_utils  import AttnProcessor
+
+import diffusers
+from diffusers import StableDiffusionXLPipeline
+from utils import PhotoMakerStableDiffusionXLPipeline
+from diffusers import DDIMScheduler
+import torch.nn.functional as F
+from utils.gradio_utils import cal_attn_mask_xl
+import copy
+import os
+from huggingface_hub import hf_hub_download
+from diffusers.utils import load_image
+from utils.utils import get_comic
+from utils.style_template import styles
+image_encoder_path = "./data/models/ip_adapter/sdxl_models/image_encoder"
+ip_ckpt = "./data/models/ip_adapter/sdxl_models/ip-adapter_sdxl_vit-h.bin"
+os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Japanese Anime"
+global models_dict
+use_va = True
+models_dict = {
+#    "Juggernaut": "RunDiffusion/Juggernaut-XL-v8",
+#    "RealVision": "SG161222/RealVisXL_V4.0" ,
+#    "SDXL":"stabilityai/stable-diffusion-xl-base-1.0" ,
+   "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y"
+}
+photomaker_path =  hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")
+MAX_SEED = np.iinfo(np.int32).max
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+def set_text_unfinished():
+    return gr.update(visible=True, value="<h3>(Not Finished) Generating ···  The intermediate results will be shown.</h3>")
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>Generation Finished</h3>")
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])
+    return image_path_list
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+################################################################################################################################################################################
+    def __init__(self, hidden_size = None, cross_attention_dim=None,id_length = 4,device = "cuda",dtype = torch.float16):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count,attn_count,cur_step,mask1024,mask4096
+        global sa32, sa64
+        global write
+        global height,width
+        if write:
+            # print(f"white:{cur_step}")
+            self.id_bank[cur_step] = [hidden_states[:self.id_length], hidden_states[self.id_length:]]
+        else:
+            encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),hidden_states[:1],self.id_bank[cur_step][1].to(self.device),hidden_states[1:]))
+        # 判断随机数是否大于0.5
+        if cur_step <5:
+            hidden_states = self.__call2__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
+        else:   # 256 1024 4096
+            random_number = random.random()
+            if cur_step <20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                # print("mask shape",mask1024.shape,mask4096.shape)
+                if not write:
+                    if hidden_states.shape[1] == (height//32) * (width//32):
+                        attention_mask = mask1024[mask1024.shape[0] // self.total_length * self.id_length:]
+                    else:
+                        attention_mask = mask4096[mask4096.shape[0] // self.total_length * self.id_length:]
+                else:
+                    # print(self.total_length,self.id_length,hidden_states.shape,(height//32) * (width//32))
+                    if hidden_states.shape[1] == (height//32) * (width//32):
+                        attention_mask = mask1024[:mask1024.shape[0] // self.total_length * self.id_length,:mask1024.shape[0] // self.total_length * self.id_length]
+                    else:
+                        attention_mask = mask4096[:mask4096.shape[0] // self.total_length * self.id_length,:mask4096.shape[0] // self.total_length * self.id_length]
+                   # print(attention_mask.shape)
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                hidden_states = self.__call1__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
+            else:
+                hidden_states = self.__call2__(attn, hidden_states,None,attention_mask,temb)
+        attn_count +=1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            mask1024,mask4096 = cal_attn_mask_xl(self.total_length,self.id_length,sa32,sa64,height,width, device=self.device, dtype= self.dtype)
+
+        return hidden_states
+    def __call1__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # print("hidden state shape",hidden_states.shape,self.id_length)
+        residual = hidden_states
+        # if encoder_hidden_states is not None:
+        #     raise Exception("not implement")
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            total_batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(total_batch_size, channel, height * width).transpose(1, 2)
+        total_batch_size,nums_token,channel = hidden_states.shape
+        img_nums = total_batch_size//2
+        hidden_states = hidden_states.view(-1,img_nums,nums_token,channel).reshape(-1,img_nums * nums_token,channel)
+
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,nums_token,channel).reshape(-1,(self.id_length+1) * nums_token,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        # print(key.shape,value.shape,query.shape,attention_mask.shape)
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        #print(query.shape,key.shape,value.shape,attention_mask.shape)
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(total_batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        # if input_ndim == 4:
+        #     tile_hidden_states = tile_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        # if attn.residual_connection:
+        #     tile_hidden_states = tile_hidden_states + residual
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(total_batch_size, channel, height, width)
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        hidden_states = hidden_states / attn.rescale_output_factor
+        # print(hidden_states.shape)
+        return hidden_states
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, channel = (
+            hidden_states.shape
+        )
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+def set_attention_processor(unet,id_length,is_ipadapter = False):
+    global total_count
+    total_count = 0
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks") :
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length = id_length)
+                total_count +=1
+            else:    
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+    print("successsfully load paired self-attention")
+    print(f"number of the processor : {total_count}")
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = '''
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+'''
+
+
+#################################################
+title = r"""
+<h1 align="center">StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</h1>
+"""
+
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+3️⃣ Choose your preferred style template.<br>
+4️⃣ Click the <b>Submit</b> button to start customizing.
+"""
+
+article = r"""
+
+If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+---
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+
+```bibtex
+@article{Zhou2024storydiffusion,
+  title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+  author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+  year={2024}
+}
+```
+📋 **License**
+<br>
+The Contents you create are under Apache-2.0 LICENSE. The Code are under Attribution-NonCommercial 4.0 International. 
+
+📧 **Contact**
+<br>
+If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+"""
+version = r"""
+<h3 align="center">StoryDiffusion Version 0.01 (test version)</h3>
+
+<h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+<h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+<h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling."),Currently, support is only using Textual Description</h5>
+<h5 align="center">Tips: Not Ready Now! Just Test</h5>
+"""
+#################################################
+global attn_count, total_count, id_length, total_length,cur_step, cur_model_type
+global write
+global  sa32, sa64
+global height,width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+device="cuda"
+global attn_procs,unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global sd_model_path
+sd_model_path = models_dict["Unstable"]#"SG161222/RealVisXL_V4.0"
+use_safetensors= False
+### LOAD Stable Diffusion Pipeline
+pipe1 = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16, use_safetensors= use_safetensors)
+pipe1 = pipe1.to("cuda")
+pipe1.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe1.scheduler.set_timesteps(50)
+### 
+pipe2 = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+    sd_model_path, torch_dtype=torch.float16, use_safetensors=use_safetensors)
+pipe2 = pipe2.to("cuda")
+pipe2.load_photomaker_adapter(
+    os.path.dirname(photomaker_path),
+    subfolder="",
+    weight_name=os.path.basename(photomaker_path),
+    trigger_word="img"  # define the trigger word
+)
+pipe2 = pipe2.to("cuda")
+pipe2.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+pipe2.fuse_lora()
+
+######### Gradio Fuction #############
+
+def swap_to_gallery(images):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+def remove_tips():
+    return gr.update(visible=False)
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive) 
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [p.replace("{prompt}", positive) for positive in positives], n + ' ' + negative
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    elif _model_type == "Using Ref Images":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
+    else:
+        raise ValueError("Invalid model type",_model_type)
+
+
+######### Image Generation ##############
+@spaces.GPU
+def process_generation(_sd_type,_model_type,_upload_images, _num_steps,style_name, _Ip_Adapter_Strength ,_style_strength_ratio, guidance_scale, seed_,  sa32_, sa64_, id_length_,  general_prompt, negative_prompt,prompt_array,G_height,G_width,_comic_type):
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error("Please add the triger word \" img \"  behind the class word you want to customize, such as: man img or woman img")
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64,id_length,total_length,attn_procs,unet,cur_model_type,device
+    global write
+    global cur_step,attn_count
+    global height,width
+    height = G_height
+    width = G_width
+    global pipe1,pipe2
+    global sd_model_path,models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    if _model_type == "original":
+        pipe = pipe1
+        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    elif _model_type == "Photomaker":
+        pipe = pipe2
+        set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    else:
+        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+        ##### ########################
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+    cur_model_type = _sd_type+"-"+_model_type+""+str(id_length_)
+    if _model_type != "original":
+        input_id_images = []
+        for img in _upload_images:
+            print(img)
+            input_id_images.append(load_image(img))
+    prompts = prompt_array.splitlines()
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device="cuda").manual_seed(seed_)
+    sa32, sa64 =  sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    prompts = [general_prompt + "," + prompt if "[NC]" not in prompt else prompt.replace("[NC]","")  for prompt in clipped_prompts]
+    prompts = [prompt.rpartition('#')[0] if "#" in prompt else prompt for prompt in prompts]
+    print(prompts)
+    id_prompts = prompts[:id_length]
+    real_prompts = prompts[id_length:]
+    torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    setup_seed(seed_)
+    total_results = []
+    if _model_type == "original":
+        id_images = pipe(id_prompts, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
+    elif _model_type == "Photomaker":
+        id_images = pipe(id_prompts,input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale, start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images
+    else: 
+        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+    total_results = id_images + total_results
+    yield total_results
+    real_images = []
+    write = False
+    for real_prompt in real_prompts:
+        setup_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":   
+            real_images.append(pipe(real_prompt,  num_inference_steps=_num_steps, guidance_scale=guidance_scale,  height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
+        elif _model_type == "Photomaker":      
+            real_images.append(pipe(real_prompt, input_id_images=input_id_images, num_inference_steps=_num_steps, guidance_scale=guidance_scale,  start_merge_step = start_merge_step, height = height, width = width,negative_prompt = negative_prompt,generator = generator).images[0])
+        else:
+            raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {_model_type}")
+        total_results = [real_images[-1]] + total_results
+        yield total_results
+    if _comic_type != "No typesetting (default)":
+        captions= prompt_array.splitlines()
+        captions = [caption.replace("[NC]","") for caption in captions]
+        captions = [caption.split('#')[-1] if "#" in caption else caption for caption in captions]
+        from PIL import ImageFont
+        total_results = get_comic(id_images + real_images, _comic_type,captions= captions,font=ImageFont.truetype("./fonts/Inkfree.ttf", int(45))) + total_results
+    set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+    yield total_results
+
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i,part in enumerate(arr):
+        if i != len(arr)-1:
+            stringtmp += part +"\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+with gr.Blocks(css=css) as demo:
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+
+    # gr.Markdown(logo)
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Group(elem_id="main-image"):
+            # button_run = gr.Button("generate id images ! 😺", elem_id="main_button", interactive=True)
+
+            prompts = []
+            colors = []
+            # with gr.Column(visible=False) as post_sketch:
+            #     for n in range(MAX_COLORS):
+            #         if n == 0 :
+            #             with gr.Row(visible=False) as color_row[n]:
+            #                 colors.append(gr.Image(shape=(100, 100), label="background", type="pil", image_mode="RGB", width=100, height=100))
+            #                 prompts.append(gr.Textbox(label="Prompt for the background (white region)", value=""))
+            #         else:
+            #             with gr.Row(visible=False) as color_row[n]:
+            #                 colors.append(gr.Image(shape=(100, 100), label="segment "+str(n), type="pil", image_mode="RGB", width=100, height=100))
+            #                 prompts.append(gr.Textbox(label="Prompt for the segment "+str(n)))
+
+            #     get_genprompt_run = gr.Button("(2) I've finished segment labeling ! 😺", elem_id="prompt_button", interactive=True)
+
+            with gr.Column(visible=True) as gen_prompt_vis:
+                sd_type = gr.Dropdown(choices=list(models_dict.keys()), value = "Unstable",label="sd_type", info="Select pretrained model")
+                model_type = gr.Radio(["Only Using Textual Description", "Using Ref Images"], label="model_type", value = "Only Using Textual Description",  info="Control type of the Character")
+                with gr.Group(visible=False) as control_image_input:
+                    files = gr.Files(
+                                label="Drag (Select) 1 or more photos of your face",
+                                file_types=["image"],
+                            )
+                    uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
+                    with gr.Column(visible=False) as clear_button:
+                        remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+                general_prompt = gr.Textbox(value='', label="(1) Textual Description for Character", interactive=True)
+                negative_prompt = gr.Textbox(value='', label="(2) Negative_prompt", interactive=True)
+                style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                prompt_array = gr.Textbox(lines = 3,value='', label="(3) Comic Description (each line corresponds to a frame).", interactive=True)
+                with gr.Accordion("(4) Tune the hyperparameters", open=True):
+                    #sa16_ = gr.Slider(label=" (The degree of Paired Attention at 16 x 16 self-attention layers) ", minimum=0, maximum=1., value=0.3, step=0.1)
+                    sa32_ = gr.Slider(label=" (The degree of Paired Attention at 32 x 32 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
+                    sa64_ = gr.Slider(label=" (The degree of Paired Attention at 64 x 64 self-attention layers) ", minimum=0, maximum=1., value=0.7, step=0.1)
+                    id_length_ = gr.Slider(label= "Number of id images in total images" , minimum=2, maximum=4, value=2, step=1)
+                    # total_length_ = gr.Slider(label= "Number of total images", minimum=1, maximum=20, value=1, step=1)
+                    seed_ = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, value=0, step=1)
+                    num_steps = gr.Slider( 
+                        label="Number of sample steps",
+                        minimum=20,
+                        maximum=100,
+                        step=1,
+                        value=50,
+                    )
+                    G_height = gr.Slider( 
+                        label="height",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    G_width = gr.Slider( 
+                        label="width",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    comic_type = gr.Radio(["No typesetting (default)", "Four Pannel", "Classic Comic Style"], value = "Classic Comic Style", label="Typesetting Style", info="Select the typesetting style ")
+                    guidance_scale = gr.Slider(
+                        label="Guidance scale",
+                        minimum=0.1,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5,
+                    )
+                    style_strength_ratio = gr.Slider(
+                        label="Style strength of Ref Image (%)",
+                        minimum=15,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                        visible=False
+                    )
+                    Ip_Adapter_Strength = gr.Slider(
+                        label="Ip_Adapter_Strength",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=0.5,
+                        visible=False
+                    )
+                final_run_btn = gr.Button("Generate ! 😺")
+
+
+        with gr.Column():
+            out_image = gr.Gallery(label="Result", columns=2, height='auto')
+            generated_information = gr.Markdown(label="Generation Details", value="",visible=False)
+            gr.Markdown(version)
+    model_type.change(fn = change_visiale_by_model_type , inputs = model_type, outputs=[control_image_input,style_strength_ratio,Ip_Adapter_Strength])
+    files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
+    remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
+
+    final_run_btn.click(fn=set_text_unfinished, outputs = generated_information
+    ).then(process_generation, inputs=[sd_type,model_type,files, num_steps,style, Ip_Adapter_Strength,style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array,G_height,G_width,comic_type], outputs=out_image
+    ).then(fn=set_text_finished,outputs = generated_information)
+
+
+    gr.Examples(
+        examples=[
+            [1,0.5,0.5,3,"a woman img, wearing a white T-shirt, blue loose hair",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                   array2string(["wake up in the bed",
+                                "have breakfast",
+                                "is on the road, go to company",
+                                "work in the company",
+                                "Take a walk next to the company at noon",
+                                "lying in bed at night"]),
+                                "Japanese Anime",  "Using Ref Images",get_image_path_list('./examples/taylor'),768,768
+                ],
+                [0,0.5,0.5,2,"a man, wearing black jacket",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                   array2string(["wake up in the bed",
+                                "have breakfast",
+                                "is on the road, go to the company,  close look",
+                                "work in the company",
+                                "laughing happily",
+                                "lying in bed at night"
+                                ]),
+                                "Japanese Anime","Only Using Textual Description",get_image_path_list('./examples/taylor'),768,768
+                ],
+                [0,0.3,0.5,2,"a girl, wearing white shirt, black skirt, black tie, yellow hair",
+                   "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                    array2string([
+                            "at home #at home, began to go to drawing",
+                            "sitting alone on a park bench.",
+                            "reading a book on a park bench.",
+                            "[NC]A squirrel approaches, peeking over the bench. ",
+                            "look around in the park. # She looks around and enjoys the beauty of nature.",
+                            "[NC]leaf falls from the tree, landing on the sketchbook.",
+                            "picks up the leaf, examining its details closely.",
+                            "starts sketching the leaf with intricate lines.",
+                            "holds up the sketch drawing of the leaf.",
+                            "[NC]The brown squirrel appear.",
+                            "is very happy # She is very happy to see the squirrel again",
+                            "[NC]The brown squirrel takes the cracker and scampers up a tree. # She gives the squirrel cracker",
+                            "laughs and tucks the leaf into her book as a keepsake.",
+                            "ready to leave.",]),
+                    "Japanese Anime","Only Using Textual Description",get_image_path_list('./examples/taylor'),768,768
+                ]
+                ],
+        inputs=[seed_, sa32_, sa64_, id_length_,  general_prompt, negative_prompt, prompt_array,style,model_type,files,G_height,G_width],
+        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
+        # run_on_click=True,
+        label='😺 Examples 😺',
+    )
+    gr.Markdown(article)
+
+    # demo.load(None, None, None, _js=load_js)
+
+demo.launch(server_name="0.0.0.0", share = True if use_va else False)

cog.yaml

@@ -0,0 +1,23 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  gpu: true
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_version: "3.11"
+  python_packages:
+    - xformers==0.0.20
+    - torch==2.0.1
+    - torchvision==0.15.2
+    - diffusers==0.25.0
+    - transformers==4.36.2
+    - gradio==3.48.0
+    - accelerate
+    - safetensors
+    - peft
+    - Pillow==9.5.0
+  run:
+    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.6.0/pget_linux_x86_64" && chmod +x /usr/local/bin/pget
+predict: "predict.py:Predictor"

gradio_app_sdxl_specific_id_low_vram copy.py

@@ -0,0 +1,1453 @@
+from this import d
+import gradio as gr
+import numpy as np
+import torch
+import gc
+import copy
+import os
+import random
+import datetime
+from PIL import ImageFont
+from utils.gradio_utils import (
+    character_to_dict,
+    process_original_prompt,
+    get_ref_character,
+    cal_attn_mask_xl,
+    cal_attn_indice_xl_effcient_memory,
+    is_torch2_available,
+)
+
+import os
+os.environ['GPU_PLATFORM_ID'] = '0'
+os.environ['GPU_DEVICE_ID'] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+import os
+os.environ['HF_ENDPOINT']= 'https://hf-mirror.com'
+torch.backends.cudnn.enabled = True
+
+if is_torch2_available():
+    from utils.gradio_utils import AttnProcessor2_0 as AttnProcessor
+else:
+    from utils.gradio_utils import AttnProcessor
+from huggingface_hub import hf_hub_download
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+)
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+import torch.nn.functional as F
+from diffusers.utils.loading_utils import load_image
+from utils.utils import get_comic
+from utils.style_template import styles
+from utils.load_models_utils import get_models_dict, load_models
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "日本漫画风"
+global models_dict
+
+models_dict = get_models_dict()
+
+# Automatically select the device
+device = (
+    "cuda:0"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+# device = "cpu"
+
+# torch.cuda.set_device(5)
+print(f"@@device:{device}")
+
+
+# check if the file exists locally at a specified path before downloading it.
+# if the file doesn't exist, it uses `hf_hub_download` to download the file
+# and optionally move it to a specific directory. If the file already
+# exists, it simply uses the local path.
+local_dir = "data/"
+photomaker_local_path = f"{local_dir}photomaker-v1.bin"
+if not os.path.exists(photomaker_local_path):
+    photomaker_path = hf_hub_download(
+        repo_id="TencentARC/PhotoMaker",
+        filename="photomaker-v1.bin",
+        repo_type="model",
+        local_dir=local_dir,
+    )
+else:
+    photomaker_path = photomaker_local_path
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def set_text_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>正在生成中......</h3>",
+    )
+
+
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>生成完成！</h3>")
+
+
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted(
+        [os.path.join(folder_name, basename) for basename in image_basename_list]
+    )
+    return image_path_list
+
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        id_length=4,
+        device=device,
+        dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count, attn_count, cur_step, indices1024, indices4096
+        global sa32, sa64
+        global write
+        global height, width
+        global character_dict, character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, cur_character
+        if attn_count == 0 and cur_step == 0:
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        if write:
+            assert len(cur_character) == 1
+            if hidden_states.shape[1] == (height // 32) * (width // 32):
+                indices = indices1024
+            else:
+                indices = indices4096
+            # print(f"white:{cur_step}")
+            total_batch_size, nums_token, channel = hidden_states.shape
+            img_nums = total_batch_size // 2
+            hidden_states = hidden_states.reshape(-1, img_nums, nums_token, channel)
+            # print(img_nums,len(indices),hidden_states.shape,self.total_length)
+            if cur_character[0] not in self.id_bank:
+                self.id_bank[cur_character[0]] = {}
+            self.id_bank[cur_character[0]][cur_step] = [
+                hidden_states[:, img_ind, indices[img_ind], :]
+                .reshape(2, -1, channel)
+                .clone()
+                for img_ind in range(img_nums)
+            ]
+            hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            # self.id_bank[cur_step] = [hidden_states[:self.id_length].clone(), hidden_states[self.id_length:].clone()]
+        else:
+            # encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),self.id_bank[cur_step][1].to(self.device)))
+            # TODO: ADD Multipersion Control
+            encoder_arr = []
+            for character in cur_character:
+                encoder_arr = encoder_arr + [
+                    tensor.to(self.device)
+                    for tensor in self.id_bank[character][cur_step]
+                ]
+        # 判断随机数是否大于0.5
+        if cur_step < 1:
+            hidden_states = self.__call2__(
+                attn, hidden_states, None, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                if hidden_states.shape[1] == (height // 32) * (width // 32):
+                    indices = indices1024
+                else:
+                    indices = indices4096
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                if write:
+                    total_batch_size, nums_token, channel = hidden_states.shape
+                    img_nums = total_batch_size // 2
+                    hidden_states = hidden_states.reshape(
+                        -1, img_nums, nums_token, channel
+                    )
+                    encoder_arr = [
+                        hidden_states[:, img_ind, indices[img_ind], :].reshape(
+                            2, -1, channel
+                        )
+                        for img_ind in range(img_nums)
+                    ]
+                    for img_ind in range(img_nums):
+                        # print(img_nums)
+                        # assert img_nums != 1
+                        img_ind_list = [i for i in range(img_nums)]
+                        # print(img_ind_list,img_ind)
+                        img_ind_list.remove(img_ind)
+                        # print(img_ind,invert_character_index_dict[img_ind])
+                        # print(character_index_dict[invert_character_index_dict[img_ind]])
+                        # print(img_ind_list)
+                        # print(img_ind,img_ind_list)
+                        encoder_hidden_states_tmp = torch.cat(
+                            [encoder_arr[img_ind] for img_ind in img_ind_list]
+                            + [hidden_states[:, img_ind, :, :]],
+                            dim=1,
+                        )
+
+                        hidden_states[:, img_ind, :, :] = self.__call2__(
+                            attn,
+                            hidden_states[:, img_ind, :, :],
+                            encoder_hidden_states_tmp,
+                            None,
+                            temb,
+                        )
+                else:
+                    _, nums_token, channel = hidden_states.shape
+                    # img_nums = total_batch_size // 2
+                    # encoder_hidden_states = encoder_hidden_states.reshape(-1,img_nums,nums_token,channel)
+                    hidden_states = hidden_states.reshape(2, -1, nums_token, channel)
+                    # print(len(indices))
+                    # encoder_arr = [encoder_hidden_states[:,img_ind,indices[img_ind],:].reshape(2,-1,channel) for img_ind in range(img_nums)]
+                    encoder_hidden_states_tmp = torch.cat(
+                        encoder_arr + [hidden_states[:, 0, :, :]], dim=1
+                    )
+                    # print(len(encoder_arr),encoder_hidden_states_tmp.shape)
+                    hidden_states[:, 0, :, :] = self.__call2__(
+                        attn,
+                        hidden_states[:, 0, :, :],
+                        encoder_hidden_states_tmp,
+                        None,
+                        temb,
+                    )
+                hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        # else:
+        #     encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global attn_procs
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = (
+            None
+            if name.endswith("attn1.processor")
+            else unet.config.cross_attention_dim
+        )
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = """
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+"""
+
+
+def save_single_character_weights(unet, character, description, filepath):
+    """
+    保存 attention_processor 类中的 id_bank GPU Tensor 列表到指定文件中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重要保存到的文件路径。
+    """
+    weights_to_save = {}
+    weights_to_save["description"] = description
+    weights_to_save["character"] = character
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 将每个 Tensor 转到 CPU 并转为列表，以确保它可以被序列化
+            weights_to_save[attn_name] = {}
+            for step_key in attn_processor.id_bank[character].keys():
+                weights_to_save[attn_name][step_key] = [
+                    tensor.cpu()
+                    for tensor in attn_processor.id_bank[character][step_key]
+                ]
+    # 使用torch.save保存权重
+    torch.save(weights_to_save, filepath)
+
+
+def load_single_character_weights(unet, filepath):
+    """
+    从指定文件中加载权重到 attention_processor 类的 id_bank 中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重文件的路径。
+    """
+    # 使用torch.load来读取权重
+    weights_to_load = torch.load(filepath, map_location=torch.device("cpu"))
+    character = weights_to_load["character"]
+    description = weights_to_load["description"]
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 转移权重到GPU（如果GPU可用的话）并赋值给id_bank
+            attn_processor.id_bank[character] = {}
+            for step_key in weights_to_load[attn_name].keys():
+                attn_processor.id_bank[character][step_key] = [
+                    tensor.to(unet.device)
+                    for tensor in weights_to_load[attn_name][step_key]
+                ]
+
+
+def save_results(unet, img_list):
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+    folder_name = f"results/{timestamp}"
+    weight_folder_name = f"{folder_name}/weights"
+    # 创建文件夹
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+        os.makedirs(weight_folder_name)
+
+    for idx, img in enumerate(img_list):
+        file_path = os.path.join(folder_name, f"image_{idx}.png")  # 图片文件名
+        img.save(file_path)
+    global character_dict
+    # for char in character_dict:
+    #     description = character_dict[char]
+    #     save_single_character_weights(unet,char,description,os.path.join(weight_folder_name, f'{char}.pt'))
+
+
+#################################################
+title = r"""
+<h1 align="center" style="font-family: 'Comic Sans MS', 'Orbitron', sans-serif; color: #00ccff; text-shadow: 2px 2px 4px #000000; font-size: 48px;">
+🧠✨ 我的AI研学旅记 🚀🤖
+</h1>
+"""
+
+
+# title = r"""
+# <h1 align="center">我的AI研学旅记</h1>
+# """
+
+description = r"""
+"""
+# <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+# ❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+# 1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+# 2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+# 3️⃣ Choose your preferred style template.<br>
+# 4️⃣ Click the <b>Submit</b> button to start customizing.
+
+
+article = r"""
+"""
+# If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+# [![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+# ---
+# 📝 **Citation**
+# <br>
+# If our work is useful for your research, please consider citing:
+
+# ```bibtex
+# @article{Zhou2024storydiffusion,
+#   title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+#   author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+#   year={2024}
+# }
+# ```
+# 📋 **License**
+# <br>
+# Apache-2.0 LICENSE. 
+
+# 📧 **Contact**
+# <br>
+# If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+
+version = r"""
+"""
+# <h3 align="center">StoryDiffusion Version 0.02 (test version)</h3>
+
+# <h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+# <h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+# <h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling.")</h5>
+# <h5 align="center">Tips: </h4>
+
+
+
+
+
+
+#################################################
+global attn_count, total_count, id_length, total_length, cur_step, cur_model_type
+global write
+global sa32, sa64
+global height, width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+global attn_procs, unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global pipe
+global sd_model_path
+pipe = None
+sd_model_path = models_dict["Unstable"]["path"]  # "SG161222/RealVisXL_V4.0"
+single_files = models_dict["Unstable"]["single_files"]
+### LOAD Stable Diffusion Pipeline
+if single_files:
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        sd_model_path, torch_dtype=torch.float16
+    )
+else:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        sd_model_path, torch_dtype=torch.float16, use_safetensors=False
+    )
+print("pipE.device = ", device)
+pipe = pipe.to(device)
+pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler.set_timesteps(50)
+pipe.enable_vae_slicing()
+if device != "mps":
+    pipe.enable_model_cpu_offload()
+unet = pipe.unet
+cur_model_type = "Unstable" + "-" + "original"
+### Insert PairedAttention
+for name in unet.attn_processors.keys():
+    cross_attention_dim = (
+        None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+    )
+    if name.startswith("mid_block"):
+        hidden_size = unet.config.block_out_channels[-1]
+    elif name.startswith("up_blocks"):
+        block_id = int(name[len("up_blocks.")])
+        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+    elif name.startswith("down_blocks"):
+        block_id = int(name[len("down_blocks.")])
+        hidden_size = unet.config.block_out_channels[block_id]
+    if cross_attention_dim is None and (name.startswith("up_blocks")):
+        attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+        total_count += 1
+    else:
+        attn_procs[name] = AttnProcessor()
+print("successsfully load paired self-attention")
+print(f"number of the processor : {total_count}")
+unet.set_attn_processor(copy.deepcopy(attn_procs))
+global mask1024, mask4096
+mask1024, mask4096 = cal_attn_mask_xl(
+    total_length,
+    id_length,
+    sa32,
+    sa64,
+    height,
+    width,
+    device=device,
+    dtype=torch.float16,
+)
+
+######### Gradio Fuction #############
+
+
+
+
+def swap_to_gallery(images):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+
+def remove_tips():
+    return gr.update(visible=False)
+
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive)
+
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [
+        p.replace("{prompt}", positive) for positive in positives
+    ], n + " " + negative
+
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif _model_type == "Using Ref Images":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    else:
+        raise ValueError("Invalid model type", _model_type)
+
+
+def load_character_files(character_files: str):
+    if character_files == "":
+        raise gr.Error("Please set a character file!")
+    character_files_arr = character_files.splitlines()
+    primarytext = []
+    for character_file_name in character_files_arr:
+        character_file = torch.load(
+            character_file_name, map_location=torch.device("cpu")
+        )
+        primarytext.append(character_file["character"] + character_file["description"])
+    return array2string(primarytext)
+
+
+def load_character_files_on_running(unet, character_files: str):
+    if character_files == "":
+        return False
+    character_files_arr = character_files.splitlines()
+    for character_file in character_files_arr:
+        load_single_character_weights(unet, character_file)
+    return True
+
+
+######### Image Generation ##############
+def process_generation(
+    _sd_type,
+    _model_type,
+    _upload_images,
+    _num_steps,
+    style_name,
+    _Ip_Adapter_Strength,
+    _style_strength_ratio,
+    guidance_scale,
+    seed_,
+    sa32_,
+    sa64_,
+    id_length_,
+    general_prompt,
+    negative_prompt,
+    prompt_array,
+    G_height,
+    G_width,
+    _comic_type,
+    font_choice,
+    _char_files,
+):  # Corrected font_choice usage
+    if len(general_prompt.splitlines()) > 5:
+        raise gr.Error(
+            "Support for more than three characters is temporarily unavailable due to VRAM limitations, but this issue will be resolved soon."
+        )
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error(
+            'Please add the triger word " img "  behind the class word you want to customize, such as: man img or woman img'
+        )
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type
+    global write
+    global cur_step, attn_count
+    global height, width
+    height = G_height
+    width = G_width
+    global pipe
+    global sd_model_path, models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    for attn_processor in pipe.unet.attn_processors.values():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            for values in attn_processor.id_bank.values():
+                del values
+            attn_processor.id_bank = {}
+            attn_processor.id_length = id_length
+            attn_processor.total_length = id_length + 1
+    gc.collect()
+    torch.cuda.empty_cache()
+    if cur_model_type != _sd_type + "-" + _model_type:
+        # apply the style template
+        ##### load pipe
+        del pipe
+        gc.collect()
+        if device == "cuda":
+            torch.cuda.empty_cache()
+        model_info = models_dict[_sd_type]
+        model_info["model_type"] = _model_type
+        print("device = ", device)
+        pipe = load_models(model_info, device=device, photomaker_path=photomaker_path)
+        set_attention_processor(pipe.unet, id_length_, is_ipadapter=False)
+        ##### ########################
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        cur_model_type = _sd_type + "-" + _model_type
+        pipe.enable_vae_slicing()
+        if device != "mps":
+            pipe.enable_model_cpu_offload()
+    else:
+        unet = pipe.unet
+        # unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+    load_chars = load_character_files_on_running(unet, character_files=_char_files)
+
+    prompts = prompt_array.splitlines()
+    global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals
+    character_dict, character_list = character_to_dict(general_prompt)
+
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device=device).manual_seed(seed_)
+    sa32, sa64 = sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    nc_indexs = []
+    for ind, prompt in enumerate(clipped_prompts):
+        if "[NC]" in prompt:
+            nc_indexs.append(ind)
+            if ind < id_length:
+                raise gr.Error(
+                    f"The first {id_length} row is id prompts, cannot use [NC]!"
+                )
+    prompts = [
+        prompt if "[NC]" not in prompt else prompt.replace("[NC]", "")
+        for prompt in clipped_prompts
+    ]
+
+    prompts = [
+        prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+    ]
+    print(prompts)
+    # id_prompts = prompts[:id_length]
+    (
+        character_index_dict,
+        invert_character_index_dict,
+        replace_prompts,
+        ref_indexs_dict,
+        ref_totals,
+    ) = process_original_prompt(character_dict, prompts.copy(), id_length)
+    if _model_type != "original":
+        input_id_images_dict = {}
+        if len(_upload_images) != len(character_dict.keys()):
+            raise gr.Error(
+                f"You upload images({len(_upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!"
+            )
+        for ind, img in enumerate(_upload_images):
+            input_id_images_dict[character_list[ind]] = [load_image(img)]
+    print(character_dict)
+    print(character_index_dict)
+    print(invert_character_index_dict)
+    # real_prompts = prompts[id_length:]
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    # id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    # print(id_prompts)
+    setup_seed(seed_)
+    total_results = []
+    id_images = []
+    results_dict = {}
+    global cur_character
+    
+    if not load_chars:
+        for character_key in character_dict.keys():
+            cur_character = [character_key]
+            ref_indexs = ref_indexs_dict[character_key]
+            print(character_key, ref_indexs)
+            current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs]
+            print(current_prompts)
+            setup_seed(seed_)
+            generator = torch.Generator(device=device).manual_seed(seed_)
+            cur_step = 0
+            cur_positive_prompts, negative_prompt = apply_style(
+                style_name, current_prompts, negative_prompt
+            )
+            if _model_type == "original":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            elif _model_type == "Photomaker":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    input_id_images=input_id_images_dict[character_key],
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    start_merge_step=start_merge_step,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            else:
+                raise NotImplementedError(
+                    "You should choice between original and Photomaker!",
+                    f"But you choice {_model_type}",
+                )
+
+            # total_results = id_images + total_results
+            # yield total_results
+            print(id_images)
+            for ind, img in enumerate(id_images):
+                print(ref_indexs[ind])
+                results_dict[ref_indexs[ind]] = img
+            # real_images = []
+            yield [results_dict[ind] for ind in results_dict.keys()]
+    write = False
+    if not load_chars:
+        real_prompts_inds = [
+            ind for ind in range(len(prompts)) if ind not in ref_totals
+        ]
+    else:
+        real_prompts_inds = [ind for ind in range(len(prompts))]
+    print(real_prompts_inds)
+
+    for real_prompts_ind in real_prompts_inds:
+        real_prompt = replace_prompts[real_prompts_ind]
+        cur_character = get_ref_character(prompts[real_prompts_ind], character_dict)
+        print(cur_character, real_prompt)
+        setup_seed(seed_)
+        if len(cur_character) > 1 and _model_type == "Photomaker":
+            raise gr.Error(
+                "Temporarily Not Support Multiple character in Ref Image Mode!"
+            )
+        generator = torch.Generator(device=device).manual_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images[0]
+        elif _model_type == "Photomaker":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                input_id_images=(
+                    input_id_images_dict[cur_character[0]]
+                    if real_prompts_ind not in nc_indexs
+                    else input_id_images_dict[character_list[0]]
+                ),
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                start_merge_step=start_merge_step,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+                nc_flag=True if real_prompts_ind in nc_indexs else False,
+            ).images[0]
+        else:
+            raise NotImplementedError(
+                "You should choice between original and Photomaker!",
+                f"But you choice {_model_type}",
+            )
+        yield [results_dict[ind] for ind in results_dict.keys()]
+    total_results = [results_dict[ind] for ind in range(len(prompts))]
+    if _comic_type != "No typesetting (default)":
+        captions = prompt_array.splitlines()
+        captions = [caption.replace("[NC]", "") for caption in captions]
+        captions = [
+            caption.split("#")[-1] if "#" in caption else caption
+            for caption in captions
+        ]
+        font_path = os.path.join("fonts", font_choice)
+        font = ImageFont.truetype(font_path, int(45))
+        total_results = (
+            get_comic(total_results, _comic_type, captions=captions, font=font)
+            + total_results
+        )
+    save_results(pipe.unet, total_results)
+
+    yield total_results
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i, part in enumerate(arr):
+        if i != len(arr) - 1:
+            stringtmp += part + "\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+
+with gr.Blocks(css=css) as demo:
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+
+    # gr.Markdown(logo)
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Group(elem_id="main-image"):
+
+            prompts = []
+            colors = []
+
+            with gr.Column(visible=True) as gen_prompt_vis:
+                with gr.Group(visible=False):
+                    sd_type = gr.Dropdown(
+                        choices=list(models_dict.keys()),
+                        value="Unstable",
+                        label="模型类型",
+                        info="选择预训练模型",
+                    )
+                    model_type = gr.Radio(
+                        ["仅使用文本描述", "使用参考图像"],
+                        label="控制模式",
+                        value="仅使用文本描述",
+                        info="角色控制方式",
+                    )
+                with gr.Group(visible=True) as control_image_input:
+                    files = gr.Files(
+                        label="拖放或选择1张或多张面部照片",
+                        file_types=["image"],
+                    )
+                    uploaded_files = gr.Gallery(
+                        label="已上传图片",
+                        visible=False,
+                        columns=5,
+                        rows=1,
+                        height=200,
+                    )
+                    with gr.Column(visible=False) as clear_button:
+                        remove_and_reupload = gr.ClearButton(
+                            value="清除并重新上传",
+                            components=files,
+                            size="sm",
+                        )
+                        
+                general_prompt = gr.Textbox(
+                    value="",
+                    lines=2,
+                    visible=False,
+                    label="(1) 角色文本描述",
+                    interactive=True,
+                )
+                negative_prompt = gr.Textbox(
+                    value="", 
+                    label="(2) 负面提示词",
+                    visible=False, 
+                    interactive=True
+                )
+                style = gr.Dropdown(
+                    label="风格模板",
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                )
+                prompt_array = gr.Textbox(
+                    lines=1,
+                    value="",
+                    visible=False,
+                    label="(3) 漫画描述（每行对应一个画格）",
+                    interactive=True,
+                )
+                char_path = gr.Textbox(
+                    lines=2,
+                    value="",
+                    visible=False,
+                    label="(可选) 角色文件路径",
+                    interactive=True,
+                )
+                char_btn = gr.Button("加载角色文件", visible=False)
+
+                with gr.Group(visible=False):
+                    font_choice = gr.Dropdown(
+                        label="选择字体",
+                        choices=[
+                            f for f in os.listdir("fonts") if f.endswith(".ttf")
+                        ],
+                        value="Inkfree.ttf",
+                        info="选择最终幻灯片的字体",
+                        interactive=True,
+                    )
+                    sa32_ = gr.Slider(
+                        label="32x32自注意力层配对注意力强度",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    sa64_ = gr.Slider(
+                        label="64x64自注意力层配对注意力强度",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    id_length_ = gr.Slider(
+                        label="总图像中包含ID图像的数量",
+                        minimum=1,
+                        maximum=4,
+                        value=1,
+                        step=1,
+                    )
+                    with gr.Row():
+                        seed_ = gr.Slider(
+                            label="随机种子", minimum=-1, maximum=MAX_SEED, value=0, step=1
+                        )
+                        randomize_seed_btn = gr.Button("🎲", size="sm")
+                    num_steps = gr.Slider(
+                        label="采样步数",
+                        minimum=20,
+                        maximum=100,
+                        step=1,
+                        value=35,
+                    )
+                    G_height = gr.Slider(
+                        label="图像高度",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    G_width = gr.Slider(
+                        label="图像宽度",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    comic_type = gr.Radio(
+                        [
+                            "默认",
+                            "四格漫画",
+                            "经典漫画风格",
+                        ],
+                        value="四格漫画",###########################################################################
+                        label="排版风格",
+                        info="选择漫画排版风格",
+                    )
+                with gr.Group(visible=False):
+                    guidance_scale = gr.Slider(
+                        label="引导尺度",
+                        minimum=0.1,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5,
+                    )
+                style_strength_ratio = gr.Slider(
+                    label="参考图像风格强度 (%)",
+                    minimum=15,
+                    maximum=50,
+                    step=1,
+                    value=20,
+                    visible=False,
+                )
+                Ip_Adapter_Strength = gr.Slider(
+                        label="IP适配器强度",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=0.5,
+                        visible=False,
+                    )
+                final_run_btn = gr.Button("😺开始生成！😺")
+
+        with gr.Column():
+            out_image = gr.Gallery(label="生成结果", columns=2, height="auto")
+            # print(out_image,"#########################################################################################################")
+            generated_information = gr.Markdown(
+                label="生成详情", value="", visible=False
+            )
+            gr.Markdown(version)
+    model_type.change(
+        fn=change_visiale_by_model_type,
+        inputs=model_type,
+        outputs=[control_image_input, style_strength_ratio, Ip_Adapter_Strength],
+    )
+    files.upload(
+        fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files]
+    )
+    remove_and_reupload.click(
+        fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files]
+    )
+    char_btn.click(fn=load_character_files, inputs=char_path, outputs=[general_prompt])
+
+    randomize_seed_btn.click(
+        fn=lambda: random.randint(-1, MAX_SEED),
+        inputs=[],
+        outputs=seed_,
+    )
+
+    final_run_btn.click(fn=set_text_unfinished, outputs=generated_information).then(
+        process_generation,
+        inputs=[
+            sd_type,
+            model_type,
+            files,
+            num_steps,
+            style,
+            Ip_Adapter_Strength,
+            style_strength_ratio,
+            guidance_scale,
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            G_height,
+            G_width,
+            comic_type,
+            font_choice,
+            char_path,
+        ],
+        outputs=out_image,
+    ).then(fn=set_text_finished, outputs=generated_information)
+    with gr.Accordion("😺 点击选择内容 😺", open=False):
+
+        gr.Markdown("### 👦 男生视角")
+        gr.Examples(
+            examples=[
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice]a man img\n[d]Chinese Rose img\n[b]a tree img\n[c]a sun flower img\n[e]a peach blossom img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice]a portrait of a man standing alone, realistic style, natural lighting",
+                        "[b] is a tall tree in a meadow",
+                        "[Alice] stands next to the [b]",
+                        "[b] is a tall tree in a meadow, photorealistic,with # [Alice]  standing beside it",
+                        "[c] a vibrant sunflower in a sunny field, close-up, golden petals, ultra-detailed",
+                        "[d] a blooming Chinese rose in a garden, macro shot, pink petals, soft focus",
+                        "[e] a peach blossom on a branch in spring, shallow depth of field, morning light",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],# get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Bob] a man img, \n[a]a flower img\n[b]a camel img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] a portrait of a man standing alone, realistic style, natural lighting",
+                        "[a] a peach blossom flower",
+                        "[b] a Bactrian camel with a desert background",
+                        "[c] a Chinese  bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Bob] stands next to a blooming peach blossom flower, its petals glowing softly under the sunlight",
+                        "[Bob] stands next to the camel",
+                        "[Bob] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Bob] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],# get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Bob] a man img, \n[a]a sunflower img\n[b]a horse img\n[c]a riverside house img\n[d]a Chinese pavilion img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] a portrait of a man standing alone, realistic style, natural lighting",
+                        "[a] a sunflower with large golden petals facing the sun",
+                        "[b] a brown horse standing on a grassy field",
+                        "[c] a traditional riverside house with white walls and black-tiled roof",
+                        "[d] an ancient Chinese multi-tiered pavilion with curved eaves",
+                        "[Bob] stands next to a sunflower, its bright petals glowing in the sunlight",
+                        "[Bob] stands next to the horse",
+                        "[Bob] stands beside a riverside house, reflected softly in the calm water",
+                        "[Bob] stands under a traditional Chinese pavilion with curved roofs and wooden columns",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ]
+        ],
+        example_labels=[
+            "与花朵合影",
+            "示例1:桃花，骆驼，桥，门楼",
+            "示例2:向日葵，马，水乡，亭子",
+        ],
+        inputs=[
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            style,
+            model_type,
+            files,
+            G_height,
+            G_width,
+        ],
+        )
+        
+        gr.Markdown("### 👧 女生视角")
+        gr.Examples(
+            examples=[
+            
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a flower img\n[b]a camel img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a peach blossom flower",
+                        "[b] a Bactrian camel with a desert background",
+                        "[c] a Chinese bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Alice] stands next to a blooming peach blossom flower, its petals glowing softly under the sunlight",
+                        "[Alice] stands beside a  camel",
+                        "[Alice] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Alice] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a tree img\n[b]an alpaca img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a tree ",
+                        "[b] an alpaca with fluffy white wool standing on grassland",
+                        "[c] a Chinese bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Alice] stands next to a tree,  under the sunlight",
+                        "[Alice] stands next to the alpaca",
+                        "[Alice] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Alice] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a sunflower img\n[b]a white goose img\n[c]a riverside house img\n[d]a Chinese pavilion img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a sunflower with large golden petals facing the sun",
+                        "[b] a white goose with smooth feathers standing near the water",
+                        "[c] a traditional riverside house with white walls and black-tiled roof",
+                        "[d] an ancient Chinese multi-tiered pavilion with curved eaves",
+                        "[Alice] stands next to a sunflower, its bright petals glowing in the sunlight",
+                        "[Alice] stands next to the white goose, its feathers clean and shining in the sun",
+                        "[Alice] stands beside a riverside house, its reflection shimmering in the gentle stream",
+                        "[Alice] stands under a Chinese pavilion with layered roofs and carved wooden pillars",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ]
+        ],
+        example_labels=[
+            "示例1:桃花，骆驼，桥，门楼",
+            "示例2:树，羊驼，桥，门楼",
+            "示例3:向日葵，鹅，水乡，亭子",
+        ],
+            inputs=[
+                seed_,
+                sa32_,
+                sa64_,
+                id_length_,
+                general_prompt,
+                negative_prompt,
+                prompt_array,
+                style,
+                model_type,
+                files,
+                G_height,
+                G_width,
+            ],
+        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
+        # run_on_click=True,
+        label="😺 请选择： 😺",
+    )
+    
+    gr.Markdown(article)
+
+
+demo.launch(server_name="0.0.0.0", share=True)

gradio_app_sdxl_specific_id_low_vram.py

@@ -0,0 +1,1454 @@
+from this import d
+import gradio as gr
+import numpy as np
+import torch
+import gc
+import copy
+import os
+import random
+import datetime
+from PIL import ImageFont
+from utils.gradio_utils import (
+    character_to_dict,
+    process_original_prompt,
+    get_ref_character,
+    cal_attn_mask_xl,
+    cal_attn_indice_xl_effcient_memory,
+    is_torch2_available,
+)
+
+import os
+os.environ['GPU_PLATFORM_ID'] = '0'
+os.environ['GPU_DEVICE_ID'] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+import os
+os.environ['HF_ENDPOINT']= 'https://hf-mirror.com'
+torch.backends.cudnn.enabled = True
+
+if is_torch2_available():
+    from utils.gradio_utils import AttnProcessor2_0 as AttnProcessor
+else:
+    from utils.gradio_utils import AttnProcessor
+from huggingface_hub import hf_hub_download
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+)
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+import torch.nn.functional as F
+from diffusers.utils.loading_utils import load_image
+from utils.utils import get_comic
+from utils.style_template import styles
+from utils.load_models_utils import get_models_dict, load_models
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "日本漫画风"
+global models_dict
+
+models_dict = get_models_dict()
+
+# Automatically select the device
+device = (
+    "cuda:0"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+# device = "cpu"
+
+# torch.cuda.set_device(5)
+print(f"@@device:{device}")
+
+
+# check if the file exists locally at a specified path before downloading it.
+# if the file doesn't exist, it uses `hf_hub_download` to download the file
+# and optionally move it to a specific directory. If the file already
+# exists, it simply uses the local path.
+local_dir = "data/"
+photomaker_local_path = f"{local_dir}photomaker-v1.bin"
+if not os.path.exists(photomaker_local_path):
+    photomaker_path = hf_hub_download(
+        repo_id="TencentARC/PhotoMaker",
+        filename="photomaker-v1.bin",
+        repo_type="model",
+        local_dir=local_dir,
+    )
+else:
+    photomaker_path = photomaker_local_path
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def set_text_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>正在生成中......</h3>",
+    )
+
+
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>生成完成！</h3>")
+
+
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted(
+        [os.path.join(folder_name, basename) for basename in image_basename_list]
+    )
+    return image_path_list
+
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        id_length=4,
+        device=device,
+        dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count, attn_count, cur_step, indices1024, indices4096
+        global sa32, sa64
+        global write
+        global height, width
+        global character_dict, character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, cur_character
+        if attn_count == 0 and cur_step == 0:
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        if write:
+            assert len(cur_character) == 1
+            if hidden_states.shape[1] == (height // 32) * (width // 32):
+                indices = indices1024
+            else:
+                indices = indices4096
+            # print(f"white:{cur_step}")
+            total_batch_size, nums_token, channel = hidden_states.shape
+            img_nums = total_batch_size // 2
+            hidden_states = hidden_states.reshape(-1, img_nums, nums_token, channel)
+            # print(img_nums,len(indices),hidden_states.shape,self.total_length)
+            if cur_character[0] not in self.id_bank:
+                self.id_bank[cur_character[0]] = {}
+            self.id_bank[cur_character[0]][cur_step] = [
+                hidden_states[:, img_ind, indices[img_ind], :]
+                .reshape(2, -1, channel)
+                .clone()
+                for img_ind in range(img_nums)
+            ]
+            hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            # self.id_bank[cur_step] = [hidden_states[:self.id_length].clone(), hidden_states[self.id_length:].clone()]
+        else:
+            # encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),self.id_bank[cur_step][1].to(self.device)))
+            # TODO: ADD Multipersion Control
+            encoder_arr = []
+            for character in cur_character:
+                encoder_arr = encoder_arr + [
+                    tensor.to(self.device)
+                    for tensor in self.id_bank[character][cur_step]
+                ]
+        # 判断随机数是否大于0.5
+        if cur_step < 1:
+            hidden_states = self.__call2__(
+                attn, hidden_states, None, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                if hidden_states.shape[1] == (height // 32) * (width // 32):
+                    indices = indices1024
+                else:
+                    indices = indices4096
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                if write:
+                    total_batch_size, nums_token, channel = hidden_states.shape
+                    img_nums = total_batch_size // 2
+                    hidden_states = hidden_states.reshape(
+                        -1, img_nums, nums_token, channel
+                    )
+                    encoder_arr = [
+                        hidden_states[:, img_ind, indices[img_ind], :].reshape(
+                            2, -1, channel
+                        )
+                        for img_ind in range(img_nums)
+                    ]
+                    for img_ind in range(img_nums):
+                        # print(img_nums)
+                        # assert img_nums != 1
+                        img_ind_list = [i for i in range(img_nums)]
+                        # print(img_ind_list,img_ind)
+                        img_ind_list.remove(img_ind)
+                        # print(img_ind,invert_character_index_dict[img_ind])
+                        # print(character_index_dict[invert_character_index_dict[img_ind]])
+                        # print(img_ind_list)
+                        # print(img_ind,img_ind_list)
+                        encoder_hidden_states_tmp = torch.cat(
+                            [encoder_arr[img_ind] for img_ind in img_ind_list]
+                            + [hidden_states[:, img_ind, :, :]],
+                            dim=1,
+                        )
+
+                        hidden_states[:, img_ind, :, :] = self.__call2__(
+                            attn,
+                            hidden_states[:, img_ind, :, :],
+                            encoder_hidden_states_tmp,
+                            None,
+                            temb,
+                        )
+                else:
+                    _, nums_token, channel = hidden_states.shape
+                    # img_nums = total_batch_size // 2
+                    # encoder_hidden_states = encoder_hidden_states.reshape(-1,img_nums,nums_token,channel)
+                    hidden_states = hidden_states.reshape(2, -1, nums_token, channel)
+                    # print(len(indices))
+                    # encoder_arr = [encoder_hidden_states[:,img_ind,indices[img_ind],:].reshape(2,-1,channel) for img_ind in range(img_nums)]
+                    encoder_hidden_states_tmp = torch.cat(
+                        encoder_arr + [hidden_states[:, 0, :, :]], dim=1
+                    )
+                    # print(len(encoder_arr),encoder_hidden_states_tmp.shape)
+                    hidden_states[:, 0, :, :] = self.__call2__(
+                        attn,
+                        hidden_states[:, 0, :, :],
+                        encoder_hidden_states_tmp,
+                        None,
+                        temb,
+                    )
+                hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        # else:
+        #     encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global attn_procs
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = (
+            None
+            if name.endswith("attn1.processor")
+            else unet.config.cross_attention_dim
+        )
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = """
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+"""
+
+
+def save_single_character_weights(unet, character, description, filepath):
+    """
+    保存 attention_processor 类中的 id_bank GPU Tensor 列表到指定文件中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重要保存到的文件路径。
+    """
+    weights_to_save = {}
+    weights_to_save["description"] = description
+    weights_to_save["character"] = character
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 将每个 Tensor 转到 CPU 并转为列表，以确保它可以被序列化
+            weights_to_save[attn_name] = {}
+            for step_key in attn_processor.id_bank[character].keys():
+                weights_to_save[attn_name][step_key] = [
+                    tensor.cpu()
+                    for tensor in attn_processor.id_bank[character][step_key]
+                ]
+    # 使用torch.save保存权重
+    torch.save(weights_to_save, filepath)
+
+
+def load_single_character_weights(unet, filepath):
+    """
+    从指定文件中加载权重到 attention_processor 类的 id_bank 中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重文件的路径。
+    """
+    # 使用torch.load来读取权重
+    weights_to_load = torch.load(filepath, map_location=torch.device("cpu"))
+    character = weights_to_load["character"]
+    description = weights_to_load["description"]
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 转移权重到GPU（如果GPU可用的话）并赋值给id_bank
+            attn_processor.id_bank[character] = {}
+            for step_key in weights_to_load[attn_name].keys():
+                attn_processor.id_bank[character][step_key] = [
+                    tensor.to(unet.device)
+                    for tensor in weights_to_load[attn_name][step_key]
+                ]
+
+
+def save_results(unet, img_list):
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+    folder_name = f"results/{timestamp}"
+    weight_folder_name = f"{folder_name}/weights"
+    # 创建文件夹
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+        os.makedirs(weight_folder_name)
+
+    for idx, img in enumerate(img_list):
+        file_path = os.path.join(folder_name, f"image_{idx}.png")  # 图片文件名
+        img.save(file_path)
+    global character_dict
+    # for char in character_dict:
+    #     description = character_dict[char]
+    #     save_single_character_weights(unet,char,description,os.path.join(weight_folder_name, f'{char}.pt'))
+
+
+#################################################
+title = r"""
+<h1 align="center" style="font-family: 'Comic Sans MS', 'Orbitron', sans-serif; color: #00ccff; text-shadow: 2px 2px 4px #000000; font-size: 48px;">
+🧠✨ 我的AI研学旅记 🚀🤖
+</h1>
+"""
+
+
+# title = r"""
+# <h1 align="center">我的AI研学旅记</h1>
+# """
+
+description = r"""
+"""
+# <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+# ❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+# 1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+# 2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+# 3️⃣ Choose your preferred style template.<br>
+# 4️⃣ Click the <b>Submit</b> button to start customizing.
+
+
+article = r"""
+"""
+# If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+# [![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+# ---
+# 📝 **Citation**
+# <br>
+# If our work is useful for your research, please consider citing:
+
+# ```bibtex
+# @article{Zhou2024storydiffusion,
+#   title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+#   author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+#   year={2024}
+# }
+# ```
+# 📋 **License**
+# <br>
+# Apache-2.0 LICENSE. 
+
+# 📧 **Contact**
+# <br>
+# If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+
+version = r"""
+"""
+# <h3 align="center">StoryDiffusion Version 0.02 (test version)</h3>
+
+# <h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+# <h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+# <h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling.")</h5>
+# <h5 align="center">Tips: </h4>
+
+
+
+
+
+
+#################################################
+global attn_count, total_count, id_length, total_length, cur_step, cur_model_type
+global write
+global sa32, sa64
+global height, width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+global attn_procs, unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global pipe
+global sd_model_path
+pipe = None
+sd_model_path = models_dict["Unstable"]["path"]  # "SG161222/RealVisXL_V4.0"
+single_files = models_dict["Unstable"]["single_files"]
+### LOAD Stable Diffusion Pipeline
+if single_files:
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        sd_model_path, torch_dtype=torch.float16
+    )
+else:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        sd_model_path, torch_dtype=torch.float16, use_safetensors=False
+    )
+print("pipE.device = ", device)
+pipe = pipe.to(device)
+pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler.set_timesteps(50)
+pipe.enable_vae_slicing()
+if device != "mps":
+    pipe.enable_model_cpu_offload()
+unet = pipe.unet
+cur_model_type = "Unstable" + "-" + "original"
+### Insert PairedAttention
+for name in unet.attn_processors.keys():
+    cross_attention_dim = (
+        None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+    )
+    if name.startswith("mid_block"):
+        hidden_size = unet.config.block_out_channels[-1]
+    elif name.startswith("up_blocks"):
+        block_id = int(name[len("up_blocks.")])
+        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+    elif name.startswith("down_blocks"):
+        block_id = int(name[len("down_blocks.")])
+        hidden_size = unet.config.block_out_channels[block_id]
+    if cross_attention_dim is None and (name.startswith("up_blocks")):
+        attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+        total_count += 1
+    else:
+        attn_procs[name] = AttnProcessor()
+print("successsfully load paired self-attention")
+print(f"number of the processor : {total_count}")
+unet.set_attn_processor(copy.deepcopy(attn_procs))
+global mask1024, mask4096
+mask1024, mask4096 = cal_attn_mask_xl(
+    total_length,
+    id_length,
+    sa32,
+    sa64,
+    height,
+    width,
+    device=device,
+    dtype=torch.float16,
+)
+
+######### Gradio Fuction #############
+
+
+def swap_to_gallery(images):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+
+def remove_tips():
+    return gr.update(visible=False)
+
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive)
+
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [
+        p.replace("{prompt}", positive) for positive in positives
+    ], n + " " + negative
+
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif _model_type == "Using Ref Images":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    else:
+        raise ValueError("Invalid model type", _model_type)
+
+
+def load_character_files(character_files: str):
+    if character_files == "":
+        raise gr.Error("Please set a character file!")
+    character_files_arr = character_files.splitlines()
+    primarytext = []
+    for character_file_name in character_files_arr:
+        character_file = torch.load(
+            character_file_name, map_location=torch.device("cpu")
+        )
+        primarytext.append(character_file["character"] + character_file["description"])
+    return array2string(primarytext)
+
+
+def load_character_files_on_running(unet, character_files: str):
+    if character_files == "":
+        return False
+    character_files_arr = character_files.splitlines()
+    for character_file in character_files_arr:
+        load_single_character_weights(unet, character_file)
+    return True
+
+
+######### Image Generation ##############
+def process_generation(
+    _sd_type,
+    _model_type,
+    _upload_images,
+    _num_steps,
+    style_name,
+    _Ip_Adapter_Strength,
+    _style_strength_ratio,
+    guidance_scale,
+    seed_,
+    sa32_,
+    sa64_,
+    id_length_,
+    general_prompt,
+    negative_prompt,
+    prompt_array,
+    G_height,
+    G_width,
+    _comic_type,
+    font_choice,
+    _char_files,
+):  # Corrected font_choice usage
+    if len(general_prompt.splitlines()) > 5:
+        raise gr.Error(
+            "Support for more than three characters is temporarily unavailable due to VRAM limitations, but this issue will be resolved soon."
+        )
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error(
+            'Please add the triger word " img "  behind the class word you want to customize, such as: man img or woman img'
+        )
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type
+    global write
+    global cur_step, attn_count
+    global height, width
+    height = G_height
+    width = G_width
+    global pipe
+    global sd_model_path, models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    for attn_processor in pipe.unet.attn_processors.values():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            for values in attn_processor.id_bank.values():
+                del values
+            attn_processor.id_bank = {}
+            attn_processor.id_length = id_length
+            attn_processor.total_length = id_length + 1
+    gc.collect()
+    torch.cuda.empty_cache()
+    if cur_model_type != _sd_type + "-" + _model_type:
+        # apply the style template
+        ##### load pipe
+        del pipe
+        gc.collect()
+        if device == "cuda":
+            torch.cuda.empty_cache()
+        model_info = models_dict[_sd_type]
+        model_info["model_type"] = _model_type
+        print("device = ", device)
+        pipe = load_models(model_info, device=device, photomaker_path=photomaker_path)
+        set_attention_processor(pipe.unet, id_length_, is_ipadapter=False)
+        ##### ########################
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        cur_model_type = _sd_type + "-" + _model_type
+        pipe.enable_vae_slicing()
+        if device != "mps":
+            pipe.enable_model_cpu_offload()
+    else:
+        unet = pipe.unet
+        # unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+    load_chars = load_character_files_on_running(unet, character_files=_char_files)
+
+    prompts = prompt_array.splitlines()
+    global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals
+    character_dict, character_list = character_to_dict(general_prompt)
+
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device=device).manual_seed(seed_)
+    sa32, sa64 = sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    nc_indexs = []
+    for ind, prompt in enumerate(clipped_prompts):
+        if "[NC]" in prompt:
+            nc_indexs.append(ind)
+            if ind < id_length:
+                raise gr.Error(
+                    f"The first {id_length} row is id prompts, cannot use [NC]!"
+                )
+    prompts = [
+        prompt if "[NC]" not in prompt else prompt.replace("[NC]", "")
+        for prompt in clipped_prompts
+    ]
+
+    prompts = [
+        prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+    ]
+    print(prompts)
+    # id_prompts = prompts[:id_length]
+    (
+        character_index_dict,
+        invert_character_index_dict,
+        replace_prompts,
+        ref_indexs_dict,
+        ref_totals,
+    ) = process_original_prompt(character_dict, prompts.copy(), id_length)
+    if _model_type != "original":
+        input_id_images_dict = {}
+        if len(_upload_images) != len(character_dict.keys()):
+            raise gr.Error(
+                f"You upload images({len(_upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!"
+            )
+        for ind, img in enumerate(_upload_images):
+            input_id_images_dict[character_list[ind]] = [load_image(img)]
+    print(character_dict)
+    print(character_index_dict)
+    print(invert_character_index_dict)
+    # real_prompts = prompts[id_length:]
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    # id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    # print(id_prompts)
+    setup_seed(seed_)
+    total_results = []
+    id_images = []
+    results_dict = {}
+    global cur_character
+    
+    if not load_chars:
+        for character_key in character_dict.keys():
+            cur_character = [character_key]
+            ref_indexs = ref_indexs_dict[character_key]
+            print(character_key, ref_indexs)
+            current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs]
+            print(current_prompts)
+            setup_seed(seed_)
+            generator = torch.Generator(device=device).manual_seed(seed_)
+            cur_step = 0
+            cur_positive_prompts, negative_prompt = apply_style(
+                style_name, current_prompts, negative_prompt
+            )
+            if _model_type == "original":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            elif _model_type == "Photomaker":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    input_id_images=input_id_images_dict[character_key],
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    start_merge_step=start_merge_step,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            else:
+                raise NotImplementedError(
+                    "You should choice between original and Photomaker!",
+                    f"But you choice {_model_type}",
+                )
+
+            # total_results = id_images + total_results
+            # yield total_results
+            print(id_images)
+            for ind, img in enumerate(id_images):
+                print(ref_indexs[ind])
+                results_dict[ref_indexs[ind]] = img
+            # real_images = []
+            yield [results_dict[ind] for ind in results_dict.keys()]
+    write = False
+    if not load_chars:
+        real_prompts_inds = [
+            ind for ind in range(len(prompts)) if ind not in ref_totals
+        ]
+    else:
+        real_prompts_inds = [ind for ind in range(len(prompts))]
+    print(real_prompts_inds)
+
+    for real_prompts_ind in real_prompts_inds:
+        real_prompt = replace_prompts[real_prompts_ind]
+        cur_character = get_ref_character(prompts[real_prompts_ind], character_dict)
+        print(cur_character, real_prompt)
+        setup_seed(seed_)
+        if len(cur_character) > 1 and _model_type == "Photomaker":
+            raise gr.Error(
+                "Temporarily Not Support Multiple character in Ref Image Mode!"
+            )
+        generator = torch.Generator(device=device).manual_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images[0]
+        elif _model_type == "Photomaker":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                input_id_images=(
+                    input_id_images_dict[cur_character[0]]
+                    if real_prompts_ind not in nc_indexs
+                    else input_id_images_dict[character_list[0]]
+                ),
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                start_merge_step=start_merge_step,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+                nc_flag=True if real_prompts_ind in nc_indexs else False,
+            ).images[0]
+        else:
+            raise NotImplementedError(
+                "You should choice between original and Photomaker!",
+                f"But you choice {_model_type}",
+            )
+        yield [results_dict[ind] for ind in results_dict.keys()]
+    total_results = [results_dict[ind] for ind in range(len(prompts))]
+    if _comic_type != "No typesetting (default)":
+        captions = prompt_array.splitlines()
+        captions = [caption.replace("[NC]", "") for caption in captions]
+        captions = [
+            caption.split("#")[-1] if "#" in caption else caption
+            for caption in captions
+        ]
+        font_path = os.path.join("fonts", font_choice)
+        font = ImageFont.truetype(font_path, int(45))
+        total_results = (
+            get_comic(total_results, _comic_type, captions=captions, font=font)
+            + total_results
+        )
+    save_results(pipe.unet, total_results)
+
+    yield total_results
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i, part in enumerate(arr):
+        if i != len(arr) - 1:
+            stringtmp += part + "\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+
+with gr.Blocks(css=css) as demo:
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+
+    # gr.Markdown(logo)
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Group(elem_id="main-image"):
+
+            prompts = []
+            colors = []
+
+            with gr.Column(visible=True) as gen_prompt_vis:
+                with gr.Group(visible=False):
+                    sd_type = gr.Dropdown(
+                        choices=list(models_dict.keys()),
+                        value="Unstable",
+                        label="模型类型",
+                        info="选择预训练模型",
+                    )
+                    model_type = gr.Radio(
+                        ["仅使用文本描述", "使用参考图像"],
+                        label="控制模式",
+                        value="仅使用文本描述",
+                        info="角色控制方式",
+                    )
+                with gr.Group(visible=True) as control_image_input:
+                    files = gr.Files(
+                        label="拖放或选择1张或多张面部照片",
+                        file_types=["image"],
+                    )
+                    uploaded_files = gr.Gallery(
+                        label="已上传图片",
+                        visible=False,
+                        columns=5,
+                        rows=1,
+                        height=200,
+                    )
+                    with gr.Column(visible=False) as clear_button:
+                        remove_and_reupload = gr.ClearButton(
+                            value="清除并重新上传",
+                            components=files,
+                            size="sm",
+                        )
+                        
+                general_prompt = gr.Textbox(
+                    value="",
+                    lines=2,
+                    visible=False,
+                    label="(1) 角色文本描述",
+                    interactive=True,
+                )
+                negative_prompt = gr.Textbox(
+                    value="", 
+                    label="(2) 负面提示词",
+                    visible=False, 
+                    interactive=True
+                )
+                style = gr.Dropdown(
+                    label="风格模板",
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                )
+                prompt_array = gr.Textbox(
+                    lines=1,
+                    value="",
+                    visible=False,
+                    label="(3) 漫画描述（每行对应一个画格）",
+                    interactive=True,
+                )
+                char_path = gr.Textbox(
+                    lines=2,
+                    value="",
+                    visible=False,
+                    label="(可选) 角色文件路径",
+                    interactive=True,
+                )
+                char_btn = gr.Button("加载角色文件", visible=False)
+
+                with gr.Group(visible=False):
+                    font_choice = gr.Dropdown(
+                        label="选择字体",
+                        choices=[
+                            f for f in os.listdir("fonts") if f.endswith(".ttf")
+                        ],
+                        value="Inkfree.ttf",
+                        info="选择最终幻灯片的字体",
+                        interactive=True,
+                    )
+                    sa32_ = gr.Slider(
+                        label="32x32自注意力层配对注意力强度",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    sa64_ = gr.Slider(
+                        label="64x64自注意力层配对注意力强度",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    id_length_ = gr.Slider(
+                        label="总图像中包含ID图像的数量",
+                        minimum=1,
+                        maximum=4,
+                        value=1,
+                        step=1,
+                    )
+                    with gr.Row():
+                        seed_ = gr.Slider(
+                            label="随机种子", minimum=-1, maximum=MAX_SEED, value=0, step=1
+                        )
+                        randomize_seed_btn = gr.Button("🎲", size="sm")
+                    num_steps = gr.Slider(
+                        label="采样步数",
+                        minimum=20,
+                        maximum=100,
+                        step=1,
+                        value=35,
+                    )
+                    G_height = gr.Slider(
+                        label="图像高度",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    G_width = gr.Slider(
+                        label="图像宽度",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    comic_type = gr.Radio(
+                        [
+                            "默认",
+                            "四格漫画",
+                            "经典漫画风格",
+                        ],
+                        value="四格漫画",###########################################################################
+                        label="排版风格",
+                        info="选择漫画排版风格",
+                    )
+                with gr.Group(visible=False):
+                    guidance_scale = gr.Slider(
+                        label="引导尺度",
+                        minimum=0.1,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5,
+                    )
+                style_strength_ratio = gr.Slider(
+                    label="参考图像风格强度 (%)",
+                    minimum=15,
+                    maximum=50,
+                    step=1,
+                    value=20,
+                    visible=False,
+                )
+                Ip_Adapter_Strength = gr.Slider(
+                        label="IP适配器强度",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=0.5,
+                        visible=False,
+                    )
+                final_run_btn = gr.Button("开始生成！😺")
+
+        with gr.Column():
+            out_image = gr.Gallery(label="生成结果", columns=2, height="auto")
+            # print(out_image,"#########################################################################################################")
+            generated_information = gr.Markdown(
+                label="生成详情", value="", visible=False
+            )
+            gr.Markdown(version)
+    model_type.change(
+        fn=change_visiale_by_model_type,
+        inputs=model_type,
+        outputs=[control_image_input, style_strength_ratio, Ip_Adapter_Strength],
+    )
+    files.upload(
+        fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files]
+    )
+    remove_and_reupload.click(
+        fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files]
+    )
+    char_btn.click(fn=load_character_files, inputs=char_path, outputs=[general_prompt])
+
+    randomize_seed_btn.click(
+        fn=lambda: random.randint(-1, MAX_SEED),
+        inputs=[],
+        outputs=seed_,
+    ) 
+
+    final_run_btn.click(fn=set_text_unfinished, outputs=generated_information).then(
+        process_generation,
+        inputs=[
+            sd_type,
+            model_type,
+            files,
+            num_steps,
+            style,
+            Ip_Adapter_Strength,
+            style_strength_ratio,
+            guidance_scale,
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            G_height,
+            G_width,
+            comic_type,
+            font_choice,
+            char_path,
+        ],
+        outputs=out_image,
+    ).then(fn=set_text_finished, outputs=generated_information)
+
+
+    with gr.Accordion("😺 点击选择内容 😺", open=False, elem_id="my_accordion"):
+
+        gr.Markdown("### 👦 男生视角")
+        gr.Examples(
+            examples=[
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Bob] a man img, \n[a]a flower img\n[b]a camel img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] a portrait of a man standing alone, realistic style, natural lighting",
+                        "[a] a peach blossom flower",
+                        "[b] a Bactrian camel with a desert background",
+                        "[c] a Chinese  bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Bob] stands next to a blooming peach blossom flower, its petals glowing softly under the sunlight",
+                        "[Bob] stands next to the camel",
+                        "[Bob] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Bob] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],# get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Bob] a man img, \n[a]a sunflower img\n[b]a horse img\n[c]a riverside house img\n[d]a Chinese pavilion img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] a portrait of a man standing alone, realistic style, natural lighting",
+                        "[a] a sunflower with large golden petals facing the sun",
+                        "[b] a brown horse standing on a grassy field",
+                        "[c] a traditional riverside house with white walls and black-tiled roof",
+                        "[d] an ancient Chinese multi-tiered pavilion with curved eaves",
+                        "[Bob] stands next to a sunflower, its bright petals glowing in the sunlight",
+                        "[Bob] stands next to the horse",
+                        "[Bob] stands beside a riverside house, reflected softly in the calm water",
+                        "[Bob] stands under a traditional Chinese pavilion with curved roofs and wooden columns",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Bob] a man img, \n[a]a tree img\n[b]an alpaca img\n[c] a bridge img\n[d] a riverside house img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] a portrait of a man standing alone, realistic style, natural lighting",
+                        "[a] a tree ",
+                        "[b] an alpaca with fluffy white wool standing on grassland",
+                        "[c] a Chinese bridge",
+                        "[d] a traditional riverside house with white walls and black-tiled roof",
+                        "[Bob] stands next to a tree, under the sunlight",
+                        "[Bob] stands next to the alpaca",
+                        "[Bob] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Bob] stands beside a riverside house, its reflection shimmering in the gentle stream",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ]
+        ],
+        example_labels=[
+            "示例1:桃花，骆驼，桥，门楼",
+            "示例2:向日葵，马，水乡，亭子",
+            "示例3:树，羊驼，桥，水乡",
+        ],
+        inputs=[
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            style,
+            model_type,
+            # files,
+            G_height,
+            G_width,
+        ],
+        )
+        
+        gr.Markdown("### 👧 女生视角")
+        gr.Examples(
+            examples=[
+            
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a flower img\n[b]a camel img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a peach blossom flower",
+                        "[b] a Bactrian camel with a desert background",
+                        "[c] a Chinese bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Alice] stands next to a blooming peach blossom flower, its petals glowing softly under the sunlight",
+                        "[Alice] stands beside a  camel",
+                        "[Alice] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Alice] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a tree img\n[b]an alpaca img\n[c] a bridge img\n[d] a gatehouse img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a tree ",
+                        "[b] an alpaca with fluffy white wool standing on grassland",
+                        "[c] a Chinese bridge",
+                        "[d] a traditional Chinese gatehouse (ornamental archway with a plaque)",
+                        "[Alice] stands next to a tree,  under the sunlight",
+                        "[Alice] stands next to the alpaca",
+                        "[Alice] stands on an ancient Chinese stone arch bridge over a quiet river",
+                        "[Alice] stands beneath a traditional Chinese gatehouse with red lanterns and wooden carvings",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                1,
+                "[Alice] a woman img, \n[a]a sunflower img\n[b]a white goose img\n[c]a riverside house img\n[d]a Chinese pavilion img",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Alice] a portrait of a woman standing alone, realistic style, natural lighting",
+                        "[a] a sunflower with large golden petals facing the sun",
+                        "[b] a white goose with smooth feathers standing near the water",
+                        "[c] a traditional riverside house with white walls and black-tiled roof",
+                        "[d] an ancient Chinese multi-tiered pavilion with curved eaves",
+                        "[Alice] stands next to a sunflower, its bright petals glowing in the sunlight",
+                        "[Alice] stands next to the white goose, its feathers clean and shining in the sun",
+                        "[Alice] stands beside a riverside house, its reflection shimmering in the gentle stream",
+                        "[Alice] stands under a Chinese pavilion with layered roofs and carved wooden pillars",
+                    ]
+                ),
+                "日本漫画风",
+                "Using Ref Images",
+                # [],  # get_image_path_list("examples/taylor"),
+                768,
+                768
+            ]
+        ],
+        example_labels=[
+            "示例1:桃花，骆驼，桥，门楼",
+            "示例2:树，羊驼，桥，门楼",
+            "示例3:向日葵，鹅，水乡，亭子",
+        ],
+            inputs=[
+                seed_,
+                sa32_,
+                sa64_,
+                id_length_,
+                general_prompt,
+                negative_prompt,
+                prompt_array,
+                style,
+                model_type,
+                # files,
+                G_height,
+                G_width,
+            ],
+        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
+        # run_on_click=True,
+        label="😺 请选择： 😺",
+    )
+    
+    gr.Markdown(article)
+
+
+demo.launch(server_name="0.0.0.0", share=True)

gradio_app_sdxl_specific_id_low_vram1.py

@@ -0,0 +1,1365 @@
+from this import d
+import gradio as gr
+import numpy as np
+import torch
+import gc
+import copy
+import os
+import random
+import datetime
+from PIL import ImageFont
+from utils.gradio_utils import (
+    character_to_dict,
+    process_original_prompt,
+    get_ref_character,
+    cal_attn_mask_xl,
+    cal_attn_indice_xl_effcient_memory,
+    is_torch2_available,
+)
+
+import os
+os.environ['GPU_PLATFORM_ID'] = '0'
+os.environ['GPU_DEVICE_ID'] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+import os
+os.environ['HF_ENDPOINT']= 'https://hf-mirror.com'
+torch.backends.cudnn.enabled = True
+
+if is_torch2_available():
+    from utils.gradio_utils import AttnProcessor2_0 as AttnProcessor
+else:
+    from utils.gradio_utils import AttnProcessor
+from huggingface_hub import hf_hub_download
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+)
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+import torch.nn.functional as F
+from diffusers.utils.loading_utils import load_image
+from utils.utils import get_comic
+from utils.style_template1 import styles
+from utils.load_models_utils import get_models_dict, load_models
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Japanese Anime"
+global models_dict
+
+models_dict = get_models_dict()
+
+# Automatically select the device
+device = (
+    "cuda:0"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+# device = "cpu"
+
+# torch.cuda.set_device(5)
+print(f"@@device:{device}")
+
+
+# check if the file exists locally at a specified path before downloading it.
+# if the file doesn't exist, it uses `hf_hub_download` to download the file
+# and optionally move it to a specific directory. If the file already
+# exists, it simply uses the local path.
+local_dir = "data/"
+photomaker_local_path = f"{local_dir}photomaker-v1.bin"
+if not os.path.exists(photomaker_local_path):
+    photomaker_path = hf_hub_download(
+        repo_id="TencentARC/PhotoMaker",
+        filename="photomaker-v1.bin",
+        repo_type="model",
+        local_dir=local_dir,
+    )
+else:
+    photomaker_path = photomaker_local_path
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def set_text_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Not Finished) Generating ···  The intermediate results will be shown.</h3>",
+    )
+
+
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>Generation Finished</h3>")
+
+
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted(
+        [os.path.join(folder_name, basename) for basename in image_basename_list]
+    )
+    return image_path_list
+
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        id_length=4,
+        device=device,
+        dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count, attn_count, cur_step, indices1024, indices4096
+        global sa32, sa64
+        global write
+        global height, width
+        global character_dict, character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, cur_character
+        if attn_count == 0 and cur_step == 0:
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        if write:
+            assert len(cur_character) == 1
+            if hidden_states.shape[1] == (height // 32) * (width // 32):
+                indices = indices1024
+            else:
+                indices = indices4096
+            # print(f"white:{cur_step}")
+            total_batch_size, nums_token, channel = hidden_states.shape
+            img_nums = total_batch_size // 2
+            hidden_states = hidden_states.reshape(-1, img_nums, nums_token, channel)
+            # print(img_nums,len(indices),hidden_states.shape,self.total_length)
+            if cur_character[0] not in self.id_bank:
+                self.id_bank[cur_character[0]] = {}
+            self.id_bank[cur_character[0]][cur_step] = [
+                hidden_states[:, img_ind, indices[img_ind], :]
+                .reshape(2, -1, channel)
+                .clone()
+                for img_ind in range(img_nums)
+            ]
+            hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            # self.id_bank[cur_step] = [hidden_states[:self.id_length].clone(), hidden_states[self.id_length:].clone()]
+        else:
+            # encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),self.id_bank[cur_step][1].to(self.device)))
+            # TODO: ADD Multipersion Control
+            encoder_arr = []
+            for character in cur_character:
+                encoder_arr = encoder_arr + [
+                    tensor.to(self.device)
+                    for tensor in self.id_bank[character][cur_step]
+                ]
+        # 判断随机数是否大于0.5
+        if cur_step < 1:
+            hidden_states = self.__call2__(
+                attn, hidden_states, None, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                if hidden_states.shape[1] == (height // 32) * (width // 32):
+                    indices = indices1024
+                else:
+                    indices = indices4096
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                if write:
+                    total_batch_size, nums_token, channel = hidden_states.shape
+                    img_nums = total_batch_size // 2
+                    hidden_states = hidden_states.reshape(
+                        -1, img_nums, nums_token, channel
+                    )
+                    encoder_arr = [
+                        hidden_states[:, img_ind, indices[img_ind], :].reshape(
+                            2, -1, channel
+                        )
+                        for img_ind in range(img_nums)
+                    ]
+                    for img_ind in range(img_nums):
+                        # print(img_nums)
+                        # assert img_nums != 1
+                        img_ind_list = [i for i in range(img_nums)]
+                        # print(img_ind_list,img_ind)
+                        img_ind_list.remove(img_ind)
+                        # print(img_ind,invert_character_index_dict[img_ind])
+                        # print(character_index_dict[invert_character_index_dict[img_ind]])
+                        # print(img_ind_list)
+                        # print(img_ind,img_ind_list)
+                        encoder_hidden_states_tmp = torch.cat(
+                            [encoder_arr[img_ind] for img_ind in img_ind_list]
+                            + [hidden_states[:, img_ind, :, :]],
+                            dim=1,
+                        )
+
+                        hidden_states[:, img_ind, :, :] = self.__call2__(
+                            attn,
+                            hidden_states[:, img_ind, :, :],
+                            encoder_hidden_states_tmp,
+                            None,
+                            temb,
+                        )
+                else:
+                    _, nums_token, channel = hidden_states.shape
+                    # img_nums = total_batch_size // 2
+                    # encoder_hidden_states = encoder_hidden_states.reshape(-1,img_nums,nums_token,channel)
+                    hidden_states = hidden_states.reshape(2, -1, nums_token, channel)
+                    # print(len(indices))
+                    # encoder_arr = [encoder_hidden_states[:,img_ind,indices[img_ind],:].reshape(2,-1,channel) for img_ind in range(img_nums)]
+                    encoder_hidden_states_tmp = torch.cat(
+                        encoder_arr + [hidden_states[:, 0, :, :]], dim=1
+                    )
+                    # print(len(encoder_arr),encoder_hidden_states_tmp.shape)
+                    hidden_states[:, 0, :, :] = self.__call2__(
+                        attn,
+                        hidden_states[:, 0, :, :],
+                        encoder_hidden_states_tmp,
+                        None,
+                        temb,
+                    )
+                hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        # else:
+        #     encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global attn_procs
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = (
+            None
+            if name.endswith("attn1.processor")
+            else unet.config.cross_attention_dim
+        )
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = """
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+"""
+
+
+def save_single_character_weights(unet, character, description, filepath):
+    """
+    保存 attention_processor 类中的 id_bank GPU Tensor 列表到指定文件中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重要保存到的文件路径。
+    """
+    weights_to_save = {}
+    weights_to_save["description"] = description
+    weights_to_save["character"] = character
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 将每个 Tensor 转到 CPU 并转为列表，以确保它可以被序列化
+            weights_to_save[attn_name] = {}
+            for step_key in attn_processor.id_bank[character].keys():
+                weights_to_save[attn_name][step_key] = [
+                    tensor.cpu()
+                    for tensor in attn_processor.id_bank[character][step_key]
+                ]
+    # 使用torch.save保存权重
+    torch.save(weights_to_save, filepath)
+
+
+def load_single_character_weights(unet, filepath):
+    """
+    从指定文件中加载权重到 attention_processor 类的 id_bank 中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重文件的路径。
+    """
+    # 使用torch.load来读取权重
+    weights_to_load = torch.load(filepath, map_location=torch.device("cpu"))
+    character = weights_to_load["character"]
+    description = weights_to_load["description"]
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 转移权重到GPU（如果GPU可用的话）并赋值给id_bank
+            attn_processor.id_bank[character] = {}
+            for step_key in weights_to_load[attn_name].keys():
+                attn_processor.id_bank[character][step_key] = [
+                    tensor.to(unet.device)
+                    for tensor in weights_to_load[attn_name][step_key]
+                ]
+
+
+def save_results(unet, img_list):
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+    folder_name = f"results/{timestamp}"
+    weight_folder_name = f"{folder_name}/weights"
+    # 创建文件夹
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+        os.makedirs(weight_folder_name)
+
+    for idx, img in enumerate(img_list):
+        file_path = os.path.join(folder_name, f"image_{idx}.png")  # 图片文件名
+        img.save(file_path)
+    global character_dict
+    # for char in character_dict:
+    #     description = character_dict[char]
+    #     save_single_character_weights(unet,char,description,os.path.join(weight_folder_name, f'{char}.pt'))
+
+
+#################################################
+title = r"""
+<h1 align="center">StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</h1>
+"""
+
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+3️��� Choose your preferred style template.<br>
+4️⃣ Click the <b>Submit</b> button to start customizing.
+"""
+
+article = r"""
+
+If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+---
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+
+```bibtex
+@article{Zhou2024storydiffusion,
+  title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+  author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+  year={2024}
+}
+```
+📋 **License**
+<br>
+Apache-2.0 LICENSE. 
+
+📧 **Contact**
+<br>
+If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+"""
+version = r"""
+<h3 align="center">StoryDiffusion Version 0.02 (test version)</h3>
+
+<h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+<h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+<h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling.")</h5>
+<h5 align="center">Tips: </h4>
+"""
+#################################################
+global attn_count, total_count, id_length, total_length, cur_step, cur_model_type
+global write
+global sa32, sa64
+global height, width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+global attn_procs, unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global pipe
+global sd_model_path
+pipe = None
+sd_model_path = models_dict["Unstable"]["path"]  # "SG161222/RealVisXL_V4.0"
+single_files = models_dict["Unstable"]["single_files"]
+### LOAD Stable Diffusion Pipeline
+if single_files:
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        sd_model_path, torch_dtype=torch.float16
+    )
+else:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        sd_model_path, torch_dtype=torch.float16, use_safetensors=False
+    )
+print("pipE.device = ", device)
+pipe = pipe.to(device)
+pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler.set_timesteps(50)
+pipe.enable_vae_slicing()
+if device != "mps":
+    pipe.enable_model_cpu_offload()
+unet = pipe.unet
+cur_model_type = "Unstable" + "-" + "original"
+### Insert PairedAttention
+for name in unet.attn_processors.keys():
+    cross_attention_dim = (
+        None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+    )
+    if name.startswith("mid_block"):
+        hidden_size = unet.config.block_out_channels[-1]
+    elif name.startswith("up_blocks"):
+        block_id = int(name[len("up_blocks.")])
+        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+    elif name.startswith("down_blocks"):
+        block_id = int(name[len("down_blocks.")])
+        hidden_size = unet.config.block_out_channels[block_id]
+    if cross_attention_dim is None and (name.startswith("up_blocks")):
+        attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+        total_count += 1
+    else:
+        attn_procs[name] = AttnProcessor()
+print("successsfully load paired self-attention")
+print(f"number of the processor : {total_count}")
+unet.set_attn_processor(copy.deepcopy(attn_procs))
+global mask1024, mask4096
+mask1024, mask4096 = cal_attn_mask_xl(
+    total_length,
+    id_length,
+    sa32,
+    sa64,
+    height,
+    width,
+    device=device,
+    dtype=torch.float16,
+)
+
+######### Gradio Fuction #############
+
+
+def swap_to_gallery(images):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+
+def remove_tips():
+    return gr.update(visible=False)
+
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive)
+
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [
+        p.replace("{prompt}", positive) for positive in positives
+    ], n + " " + negative
+
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif _model_type == "Using Ref Images":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    else:
+        raise ValueError("Invalid model type", _model_type)
+
+
+def load_character_files(character_files: str):
+    if character_files == "":
+        raise gr.Error("Please set a character file!")
+    character_files_arr = character_files.splitlines()
+    primarytext = []
+    for character_file_name in character_files_arr:
+        character_file = torch.load(
+            character_file_name, map_location=torch.device("cpu")
+        )
+        primarytext.append(character_file["character"] + character_file["description"])
+    return array2string(primarytext)
+
+
+def load_character_files_on_running(unet, character_files: str):
+    if character_files == "":
+        return False
+    character_files_arr = character_files.splitlines()
+    for character_file in character_files_arr:
+        load_single_character_weights(unet, character_file)
+    return True
+
+
+######### Image Generation ##############
+def process_generation(
+    _sd_type,
+    _model_type,
+    _upload_images,
+    _num_steps,
+    style_name,
+    _Ip_Adapter_Strength,
+    _style_strength_ratio,
+    guidance_scale,
+    seed_,
+    sa32_,
+    sa64_,
+    id_length_,
+    general_prompt,
+    negative_prompt,
+    prompt_array,
+    G_height,
+    G_width,
+    _comic_type,
+    font_choice,
+    _char_files,
+):  # Corrected font_choice usage
+    if len(general_prompt.splitlines()) > 5:
+        raise gr.Error(
+            "Support for more than three characters is temporarily unavailable due to VRAM limitations, but this issue will be resolved soon."
+        )
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error(
+            'Please add the triger word " img "  behind the class word you want to customize, such as: man img or woman img'
+        )
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type
+    global write
+    global cur_step, attn_count
+    global height, width
+    height = G_height
+    width = G_width
+    global pipe
+    global sd_model_path, models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    for attn_processor in pipe.unet.attn_processors.values():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            for values in attn_processor.id_bank.values():
+                del values
+            attn_processor.id_bank = {}
+            attn_processor.id_length = id_length
+            attn_processor.total_length = id_length + 1
+    gc.collect()
+    torch.cuda.empty_cache()
+    if cur_model_type != _sd_type + "-" + _model_type:
+        # apply the style template
+        ##### load pipe
+        del pipe
+        gc.collect()
+        if device == "cuda":
+            torch.cuda.empty_cache()
+        model_info = models_dict[_sd_type]
+        model_info["model_type"] = _model_type
+        print("device = ", device)
+        pipe = load_models(model_info, device=device, photomaker_path=photomaker_path)
+        set_attention_processor(pipe.unet, id_length_, is_ipadapter=False)
+        ##### ########################
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        cur_model_type = _sd_type + "-" + _model_type
+        pipe.enable_vae_slicing()
+        if device != "mps":
+            pipe.enable_model_cpu_offload()
+    else:
+        unet = pipe.unet
+        # unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+    load_chars = load_character_files_on_running(unet, character_files=_char_files)
+
+    prompts = prompt_array.splitlines()
+    global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals
+    character_dict, character_list = character_to_dict(general_prompt)
+
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device=device).manual_seed(seed_)
+    sa32, sa64 = sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    nc_indexs = []
+    for ind, prompt in enumerate(clipped_prompts):
+        if "[NC]" in prompt:
+            nc_indexs.append(ind)
+            if ind < id_length:
+                raise gr.Error(
+                    f"The first {id_length} row is id prompts, cannot use [NC]!"
+                )
+    prompts = [
+        prompt if "[NC]" not in prompt else prompt.replace("[NC]", "")
+        for prompt in clipped_prompts
+    ]
+
+    prompts = [
+        prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+    ]
+    print(prompts)
+    # id_prompts = prompts[:id_length]
+    (
+        character_index_dict,
+        invert_character_index_dict,
+        replace_prompts,
+        ref_indexs_dict,
+        ref_totals,
+    ) = process_original_prompt(character_dict, prompts.copy(), id_length)
+    if _model_type != "original":
+        input_id_images_dict = {}
+        if len(_upload_images) != len(character_dict.keys()):
+            raise gr.Error(
+                f"You upload images({len(_upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!"
+            )
+        for ind, img in enumerate(_upload_images):
+            input_id_images_dict[character_list[ind]] = [load_image(img)]
+    print(character_dict)
+    print(character_index_dict)
+    print(invert_character_index_dict)
+    # real_prompts = prompts[id_length:]
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    # id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    # print(id_prompts)
+    setup_seed(seed_)
+    total_results = []
+    id_images = []
+    results_dict = {}
+    global cur_character
+    
+    if not load_chars:
+        for character_key in character_dict.keys():
+            cur_character = [character_key]
+            ref_indexs = ref_indexs_dict[character_key]
+            print(character_key, ref_indexs)
+            current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs]
+            print(current_prompts)
+            setup_seed(seed_)
+            generator = torch.Generator(device=device).manual_seed(seed_)
+            cur_step = 0
+            cur_positive_prompts, negative_prompt = apply_style(
+                style_name, current_prompts, negative_prompt
+            )
+            if _model_type == "original":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            elif _model_type == "Photomaker":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    input_id_images=input_id_images_dict[character_key],
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    start_merge_step=start_merge_step,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            else:
+                raise NotImplementedError(
+                    "You should choice between original and Photomaker!",
+                    f"But you choice {_model_type}",
+                )
+
+            # total_results = id_images + total_results
+            # yield total_results
+            print(id_images)
+            for ind, img in enumerate(id_images):
+                print(ref_indexs[ind])
+                results_dict[ref_indexs[ind]] = img
+            # real_images = []
+            yield [results_dict[ind] for ind in results_dict.keys()]
+    write = False
+    if not load_chars:
+        real_prompts_inds = [
+            ind for ind in range(len(prompts)) if ind not in ref_totals
+        ]
+    else:
+        real_prompts_inds = [ind for ind in range(len(prompts))]
+    print(real_prompts_inds)
+
+    for real_prompts_ind in real_prompts_inds:
+        real_prompt = replace_prompts[real_prompts_ind]
+        cur_character = get_ref_character(prompts[real_prompts_ind], character_dict)
+        print(cur_character, real_prompt)
+        setup_seed(seed_)
+        if len(cur_character) > 2 and _model_type == "Photomaker":
+            raise gr.Error(
+                "Temporarily Not Support Multiple character in Ref Image Mode!"
+            )
+        generator = torch.Generator(device=device).manual_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images[0]
+        elif _model_type == "Photomaker":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                input_id_images=(
+                    input_id_images_dict[cur_character[0]]
+                    if real_prompts_ind not in nc_indexs
+                    else input_id_images_dict[character_list[0]]
+                ),
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                start_merge_step=start_merge_step,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+                nc_flag=True if real_prompts_ind in nc_indexs else False,
+            ).images[0]
+        else:
+            raise NotImplementedError(
+                "You should choice between original and Photomaker!",
+                f"But you choice {_model_type}",
+            )
+        yield [results_dict[ind] for ind in results_dict.keys()]
+    total_results = [results_dict[ind] for ind in range(len(prompts))]
+    if _comic_type != "No typesetting (default)":
+        captions = prompt_array.splitlines()
+        captions = [caption.replace("[NC]", "") for caption in captions]
+        captions = [
+            caption.split("#")[-1] if "#" in caption else caption
+            for caption in captions
+        ]
+        font_path = os.path.join("fonts", font_choice)
+        font = ImageFont.truetype(font_path, int(45))
+        total_results = (
+            get_comic(total_results, _comic_type, captions=captions, font=font)
+            + total_results
+        )
+    save_results(pipe.unet, total_results)
+
+    yield total_results
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i, part in enumerate(arr):
+        if i != len(arr) - 1:
+            stringtmp += part + "\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+
+with gr.Blocks(css=css) as demo:
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+
+    # gr.Markdown(logo)
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Group(elem_id="main-image"):
+
+            prompts = []
+            colors = []
+
+            with gr.Column(visible=True) as gen_prompt_vis:
+                sd_type = gr.Dropdown(
+                    choices=list(models_dict.keys()),
+                    value="Unstable",
+                    label="sd_type",
+                    info="Select pretrained model",
+                )
+                model_type = gr.Radio(
+                    ["Only Using Textual Description", "Using Ref Images"],
+                    label="model_type",
+                    value="Only Using Textual Description",
+                    info="Control type of the Character",
+                )
+                with gr.Group(visible=False) as control_image_input:
+                    files = gr.Files(
+                        label="Drag (Select) 1 or more photos of your face",
+                        file_types=["image"],
+                    )
+                    uploaded_files = gr.Gallery(
+                        label="Your images",
+                        visible=False,
+                        columns=5,
+                        rows=1,
+                        height=200,
+                    )
+                    with gr.Column(visible=False) as clear_button:
+                        remove_and_reupload = gr.ClearButton(
+                            value="Remove and upload new ones",
+                            components=files,
+                            size="sm",
+                        )
+                general_prompt = gr.Textbox(
+                    value="",
+                    lines=2,
+                    label="(1) Textual Description for Character",
+                    interactive=True,
+                )
+                negative_prompt = gr.Textbox(
+                    value="", label="(2) Negative_prompt", interactive=True
+                )
+                style = gr.Dropdown(
+                    label="Style template",
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                )
+                prompt_array = gr.Textbox(
+                    lines=1,
+                    value="",
+                    label="(3) Comic Description (each line corresponds to a frame).",
+                    interactive=True,
+                )
+                char_path = gr.Textbox(
+                    lines=2,
+                    value="",
+                    visible=False,
+                    label="(Optional) Character files",
+                    interactive=True,
+                )
+                char_btn = gr.Button("Load Character files", visible=False)
+                with gr.Accordion("(4) Tune the hyperparameters", open=True):
+                    font_choice = gr.Dropdown(
+                        label="Select Font",
+                        choices=[
+                            f for f in os.listdir("fonts") if f.endswith(".ttf")
+                        ],
+                        value="Inkfree.ttf",
+                        info="Select font for the final slide.",
+                        interactive=True,
+                    )
+                    sa32_ = gr.Slider(
+                        label=" (The degree of Paired Attention at 32 x 32 self-attention layers) ",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    sa64_ = gr.Slider(
+                        label=" (The degree of Paired Attention at 64 x 64 self-attention layers) ",
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.1,
+                    )
+                    id_length_ = gr.Slider(
+                        label="Number of id images in total images",
+                        minimum=1,
+                        maximum=4,
+                        value=1,
+                        step=1,
+                    )
+                    with gr.Row():
+                        seed_ = gr.Slider(
+                            label="Seed", minimum=-1, maximum=MAX_SEED, value=0, step=1
+                        )
+                        randomize_seed_btn = gr.Button("🎲", size="sm")
+                    num_steps = gr.Slider(
+                        label="Number of sample steps",
+                        minimum=20,
+                        maximum=100,
+                        step=1,
+                        value=35,
+                    )
+                    G_height = gr.Slider(
+                        label="height",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    G_width = gr.Slider(
+                        label="width",
+                        minimum=256,
+                        maximum=1024,
+                        step=32,
+                        value=768,
+                    )
+                    comic_type = gr.Radio(
+                        [
+                            "No typesetting (default)",
+                            "Four Pannel",
+                            "Classic Comic Style",
+                        ],
+                        value="Classic Comic Style",
+                        label="Typesetting Style",
+                        info="Select the typesetting style ",
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance scale",
+                        minimum=0.1,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5,
+                    )
+                    style_strength_ratio = gr.Slider(
+                        label="Style strength of Ref Image (%)",
+                        minimum=15,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                        visible=False,
+                    )
+                    Ip_Adapter_Strength = gr.Slider(
+                        label="Ip_Adapter_Strength",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=0.5,
+                        visible=False,
+                    )
+                final_run_btn = gr.Button("Generate ! 😺")
+
+        with gr.Column():
+            out_image = gr.Gallery(label="Result", columns=2, height="auto")
+            generated_information = gr.Markdown(
+                label="Generation Details", value="", visible=False
+            )
+            gr.Markdown(version)
+    model_type.change(
+        fn=change_visiale_by_model_type,
+        inputs=model_type,
+        outputs=[control_image_input, style_strength_ratio, Ip_Adapter_Strength],
+    )
+    files.upload(
+        fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files]
+    )
+    remove_and_reupload.click(
+        fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files]
+    )
+    char_btn.click(fn=load_character_files, inputs=char_path, outputs=[general_prompt])
+
+    randomize_seed_btn.click(
+        fn=lambda: random.randint(-1, MAX_SEED),
+        inputs=[],
+        outputs=seed_,
+    )
+
+    final_run_btn.click(fn=set_text_unfinished, outputs=generated_information).then(
+        process_generation,
+        inputs=[
+            sd_type,
+            model_type,
+            files,
+            num_steps,
+            style,
+            Ip_Adapter_Strength,
+            style_strength_ratio,
+            guidance_scale,
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            G_height,
+            G_width,
+            comic_type,
+            font_choice,
+            char_path,
+        ],
+        outputs=out_image,
+    ).then(fn=set_text_finished, outputs=generated_information)
+
+    gr.Examples(
+        examples=[
+            [
+                0,
+                0.5,
+                0.5,
+                2,
+                "[Bob] A man, wearing a black suit\n[Alice]a woman, wearing a white shirt",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] at home, read new paper #at home, The newspaper says there is a treasure house in the forest.",
+                        "[Bob] on the road, near the forest",
+                        "[Alice] is make a call at home # [Bob] invited [Alice] to join him on an adventure.",
+                        "[NC]A tiger appeared in the forest, at night ",
+                        "[NC] The car on the road, near the forest #They drives to the forest in search of treasure.",
+                        "[Bob] very frightened, open mouth, in the forest, at night",
+                        "[Alice] very frightened, open mouth, in the forest, at night",
+                        "[Bob]  and [Alice] running very fast, in the forest, at night",
+                        "[NC] A house in the forest, at night #Suddenly, They discovers the treasure house!",
+                        "[Bob]  and [Alice]  in the house filled with  treasure, laughing, at night #He is overjoyed inside the house.",
+                    ]
+                ),
+                "Comic book",
+                "Only Using Textual Description",
+                get_image_path_list("examples/taylor"),
+                768,
+                768,
+            ],
+            [
+                0,
+                0.5,
+                0.5,
+                2,
+                "[Bob] A man img, wearing a black suit\n[Alice]a woman img, wearing a white shirt",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob] at home, read new paper #at home, The newspaper says there is a treasure house in the forest.",
+                        "[Bob] on the road, near the forest",
+                        "[Alice] is make a call at home # [Bob] invited [Alice] to join him on an adventure.",
+                        "[NC] The car on the road, near the forest #They drives to the forest in search of treasure.",
+                        "[NC]A tiger appeared in the forest, at night ",
+                        "[Bob] very frightened, open mouth, in the forest, at night",
+                        "[Alice] very frightened, open mouth, in the forest, at night",
+                        "[Bob]  running very fast, in the forest, at night",
+                        "[NC] A house in the forest, at night #Suddenly, They discovers the treasure house!",
+                        "[Bob]  in the house filled with  treasure, laughing, at night #They are overjoyed inside the house.",
+                    ]
+                ),
+                "Comic book",
+                "Using Ref Images",
+                get_image_path_list("examples/twoperson"),
+                1024,
+                1024,
+            ],
+            [
+                1,
+                0.5,
+                0.5,
+                3,
+                "[Taylor]a woman img, wearing a white T-shirt, blue loose hair",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Taylor]wake up in the bed",
+                        "[Taylor]have breakfast",
+                        "[Taylor]is on the road, go to company",
+                        "[Taylor]work in the company",
+                        "[Taylor]Take a walk next to the company at noon",
+                        "[Taylor]lying in bed at night",
+                    ]
+                ),
+                "Japanese Anime",
+                "Using Ref Images",
+                get_image_path_list("examples/taylor"),
+                768,
+                768,
+            ],
+            [
+                0,
+                0.5,
+                0.5,
+                3,
+                "[Bob]a man, wearing black jacket",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Bob]wake up in the bed",
+                        "[Bob]have breakfast",
+                        "[Bob]is on the road, go to the company,  close look",
+                        "[Bob]work in the company",
+                        "[Bob]laughing happily",
+                        "[Bob]lying in bed at night",
+                    ]
+                ),
+                "Japanese Anime",
+                "Only Using Textual Description",
+                get_image_path_list("examples/taylor"),
+                768,
+                768,
+            ],
+            [
+                0,
+                0.3,
+                0.5,
+                3,
+                "[Kitty]a girl, wearing white shirt, black skirt, black tie, yellow hair",
+                "bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+                array2string(
+                    [
+                        "[Kitty]at home #at home, began to go to drawing",
+                        "[Kitty]sitting alone on a park bench.",
+                        "[Kitty]reading a book on a park bench.",
+                        "[NC]A squirrel approaches, peeking over the bench. ",
+                        "[Kitty]look around in the park. # She looks around and enjoys the beauty of nature.",
+                        "[NC]leaf falls from the tree, landing on the sketchbook.",
+                        "[Kitty]picks up the leaf, examining its details closely.",
+                        "[NC]The brown squirrel appear.",
+                        "[Kitty]is very happy # She is very happy to see the squirrel again",
+                        "[NC]The brown squirrel takes the cracker and scampers up a tree. # She gives the squirrel cracker",
+                    ]
+                ),
+                "Japanese Anime",
+                "Only Using Textual Description",
+                get_image_path_list("examples/taylor"),
+                768,
+                768,
+            ],
+        ],
+        inputs=[
+            seed_,
+            sa32_,
+            sa64_,
+            id_length_,
+            general_prompt,
+            negative_prompt,
+            prompt_array,
+            style,
+            model_type,
+            files,
+            G_height,
+            G_width,
+        ],
+        # outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],
+        # run_on_click=True,
+        label="😺 Examples 😺",
+    )
+    gr.Markdown(article)
+
+
+demo.launch(server_name="0.0.0.0", share=True)

gradio_app_sdxl_specific_id_low_vram1231231.py

@@ -0,0 +1,1187 @@
+from this import d
+import gradio as gr
+import numpy as np
+import torch
+import gc
+import copy
+import os
+import random
+import datetime
+from PIL import ImageFont
+from utils.gradio_utils import (
+    character_to_dict,
+    process_original_prompt,
+    get_ref_character,
+    cal_attn_mask_xl,
+    cal_attn_indice_xl_effcient_memory,
+    is_torch2_available,
+)
+
+import os
+os.environ['GPU_PLATFORM_ID'] = '0'
+os.environ['GPU_DEVICE_ID'] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+import os
+os.environ['HF_ENDPOINT']= 'https://hf-mirror.com'
+torch.backends.cudnn.enabled = True
+
+if is_torch2_available():
+    from utils.gradio_utils import AttnProcessor2_0 as AttnProcessor
+else:
+    from utils.gradio_utils import AttnProcessor
+from huggingface_hub import hf_hub_download
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+)
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+import torch.nn.functional as F
+from diffusers.utils.loading_utils import load_image
+from utils.utils import get_comic
+from utils.style_template import styles
+from utils.load_models_utils import get_models_dict, load_models
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '5'
+
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Japanese Anime"
+global models_dict
+
+models_dict = get_models_dict()
+
+# Automatically select the device
+device = (
+    "cuda:0"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+# device = "cpu"
+
+# torch.cuda.set_device(5)
+print(f"@@device:{device}")
+
+
+# check if the file exists locally at a specified path before downloading it.
+# if the file doesn't exist, it uses `hf_hub_download` to download the file
+# and optionally move it to a specific directory. If the file already
+# exists, it simply uses the local path.
+local_dir = "data/"
+photomaker_local_path = f"{local_dir}photomaker-v1.bin"
+if not os.path.exists(photomaker_local_path):
+    photomaker_path = hf_hub_download(
+        repo_id="TencentARC/PhotoMaker",
+        filename="photomaker-v1.bin",
+        repo_type="model",
+        local_dir=local_dir,
+    )
+else:
+    photomaker_path = photomaker_local_path
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def set_text_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Not Finished) Generating ···  The intermediate results will be shown.</h3>",
+    )
+
+
+def set_text_finished():
+    return gr.update(visible=True, value="<h3>Generation Finished</h3>")
+
+
+#################################################
+def get_image_path_list(folder_name):
+    image_basename_list = os.listdir(folder_name)
+    image_path_list = sorted(
+        [os.path.join(folder_name, basename) for basename in image_basename_list]
+    )
+    return image_path_list
+
+
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        id_length=4,
+        device=device,
+        dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count, attn_count, cur_step, indices1024, indices4096
+        global sa32, sa64
+        global write
+        global height, width
+        global character_dict, character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, cur_character
+        if attn_count == 0 and cur_step == 0:
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        if write:
+            assert len(cur_character) == 1
+            if hidden_states.shape[1] == (height // 32) * (width // 32):
+                indices = indices1024
+            else:
+                indices = indices4096
+            # print(f"white:{cur_step}")
+            total_batch_size, nums_token, channel = hidden_states.shape
+            img_nums = total_batch_size // 2
+            hidden_states = hidden_states.reshape(-1, img_nums, nums_token, channel)
+            # print(img_nums,len(indices),hidden_states.shape,self.total_length)
+            if cur_character[0] not in self.id_bank:
+                self.id_bank[cur_character[0]] = {}
+            self.id_bank[cur_character[0]][cur_step] = [
+                hidden_states[:, img_ind, indices[img_ind], :]
+                .reshape(2, -1, channel)
+                .clone()
+                for img_ind in range(img_nums)
+            ]
+            hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            # self.id_bank[cur_step] = [hidden_states[:self.id_length].clone(), hidden_states[self.id_length:].clone()]
+        else:
+            # encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),self.id_bank[cur_step][1].to(self.device)))
+            # TODO: ADD Multipersion Control
+            encoder_arr = []
+            for character in cur_character:
+                encoder_arr = encoder_arr + [
+                    tensor.to(self.device)
+                    for tensor in self.id_bank[character][cur_step]
+                ]
+        # 判断随机数是否大于0.5
+        if cur_step < 1:
+            hidden_states = self.__call2__(
+                attn, hidden_states, None, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                if hidden_states.shape[1] == (height // 32) * (width // 32):
+                    indices = indices1024
+                else:
+                    indices = indices4096
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                if write:
+                    total_batch_size, nums_token, channel = hidden_states.shape
+                    img_nums = total_batch_size // 2
+                    hidden_states = hidden_states.reshape(
+                        -1, img_nums, nums_token, channel
+                    )
+                    encoder_arr = [
+                        hidden_states[:, img_ind, indices[img_ind], :].reshape(
+                            2, -1, channel
+                        )
+                        for img_ind in range(img_nums)
+                    ]
+                    for img_ind in range(img_nums):
+                        # print(img_nums)
+                        # assert img_nums != 1
+                        img_ind_list = [i for i in range(img_nums)]
+                        # print(img_ind_list,img_ind)
+                        img_ind_list.remove(img_ind)
+                        # print(img_ind,invert_character_index_dict[img_ind])
+                        # print(character_index_dict[invert_character_index_dict[img_ind]])
+                        # print(img_ind_list)
+                        # print(img_ind,img_ind_list)
+                        encoder_hidden_states_tmp = torch.cat(
+                            [encoder_arr[img_ind] for img_ind in img_ind_list]
+                            + [hidden_states[:, img_ind, :, :]],
+                            dim=1,
+                        )
+
+                        hidden_states[:, img_ind, :, :] = self.__call2__(
+                            attn,
+                            hidden_states[:, img_ind, :, :],
+                            encoder_hidden_states_tmp,
+                            None,
+                            temb,
+                        )
+                else:
+                    _, nums_token, channel = hidden_states.shape
+                    # img_nums = total_batch_size // 2
+                    # encoder_hidden_states = encoder_hidden_states.reshape(-1,img_nums,nums_token,channel)
+                    hidden_states = hidden_states.reshape(2, -1, nums_token, channel)
+                    # print(len(indices))
+                    # encoder_arr = [encoder_hidden_states[:,img_ind,indices[img_ind],:].reshape(2,-1,channel) for img_ind in range(img_nums)]
+                    encoder_hidden_states_tmp = torch.cat(
+                        encoder_arr + [hidden_states[:, 0, :, :]], dim=1
+                    )
+                    # print(len(encoder_arr),encoder_hidden_states_tmp.shape)
+                    hidden_states[:, 0, :, :] = self.__call2__(
+                        attn,
+                        hidden_states[:, 0, :, :],
+                        encoder_hidden_states_tmp,
+                        None,
+                        temb,
+                    )
+                hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        # else:
+        #     encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global attn_procs
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = (
+            None
+            if name.endswith("attn1.processor")
+            else unet.config.cross_attention_dim
+        )
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+
+#################################################
+#################################################
+canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+
+get_js_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+css = """
+#color-bg{display:flex;justify-content: center;align-items: center;}
+.color-bg-item{width: 100%; height: 32px}
+#main_button{width:100%}
+<style>
+"""
+
+
+def save_single_character_weights(unet, character, description, filepath):
+    """
+    保存 attention_processor 类中的 id_bank GPU Tensor 列表到指定文件中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重要保存到的文件路径。
+    """
+    weights_to_save = {}
+    weights_to_save["description"] = description
+    weights_to_save["character"] = character
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 将每个 Tensor 转到 CPU 并转为列表，以确保它可以被序列化
+            weights_to_save[attn_name] = {}
+            for step_key in attn_processor.id_bank[character].keys():
+                weights_to_save[attn_name][step_key] = [
+                    tensor.cpu()
+                    for tensor in attn_processor.id_bank[character][step_key]
+                ]
+    # 使用torch.save保存权重
+    torch.save(weights_to_save, filepath)
+
+
+def load_single_character_weights(unet, filepath):
+    """
+    从指定文件中加载权重到 attention_processor 类的 id_bank 中。
+    参数:
+    - model: 包含 attention_processor 类实例的模型。
+    - filepath: 权重文件的路径。
+    """
+    # 使用torch.load来读取权重
+    weights_to_load = torch.load(filepath, map_location=torch.device("cpu"))
+    character = weights_to_load["character"]
+    description = weights_to_load["description"]
+    for attn_name, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            # 转移权重到GPU（如果GPU可用的话）并赋值给id_bank
+            attn_processor.id_bank[character] = {}
+            for step_key in weights_to_load[attn_name].keys():
+                attn_processor.id_bank[character][step_key] = [
+                    tensor.to(unet.device)
+                    for tensor in weights_to_load[attn_name][step_key]
+                ]
+
+
+def save_results(unet, img_list):
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+    folder_name = f"results/{timestamp}"
+    weight_folder_name = f"{folder_name}/weights"
+    # 创建文件夹
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+        os.makedirs(weight_folder_name)
+
+    for idx, img in enumerate(img_list):
+        file_path = os.path.join(folder_name, f"image_{idx}.png")  # 图片文件名
+        img.save(file_path)
+    global character_dict
+    # for char in character_dict:
+    #     description = character_dict[char]
+    #     save_single_character_weights(unet,char,description,os.path.join(weight_folder_name, f'{char}.pt'))
+
+
+#################################################
+title = r"""
+<h1 align="center">StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</h1>
+"""
+
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'><b>StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation</b></a>.<br>
+❗️❗️❗️[<b>Important</b>] Personalization steps:<br>
+1️⃣ Enter a Textual Description for Character, if you add the Ref-Image, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>
+2️⃣ Enter the prompt array, each line corrsponds to one generated image.<br>
+3️�� Choose your preferred style template.<br>
+4️⃣ Click the <b>Submit</b> button to start customizing.
+"""
+
+article = r"""
+
+If StoryDiffusion is helpful, please help to ⭐ the <a href='https://github.com/HVision-NKU/StoryDiffusion' target='_blank'>Github Repo</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/HVision-NKU/StoryDiffusion?style=social)](https://github.com/HVision-NKU/StoryDiffusion)
+---
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+
+```bibtex
+@article{Zhou2024storydiffusion,
+  title={StoryDiffusion: Consistent Self-Attention for Long-Range Image and Video Generation},
+  author={Zhou, Yupeng and Zhou, Daquan and Cheng, Ming-Ming and Feng, Jiashi and Hou, Qibin},
+  year={2024}
+}
+```
+📋 **License**
+<br>
+Apache-2.0 LICENSE. 
+
+📧 **Contact**
+<br>
+If you have any questions, please feel free to reach me out at <b>ypzhousdu@gmail.com</b>.
+"""
+version = r"""
+<h3 align="center">StoryDiffusion Version 0.02 (test version)</h3>
+
+<h5 >1. Support image ref image. (Cartoon Ref image is not support now)</h5>
+<h5 >2. Support Typesetting Style and Captioning.(By default, the prompt is used as the caption for each image. If you need to change the caption, add a # at the end of each line. Only the part after the # will be added as a caption to the image.)</h5>
+<h5 >3. [NC]symbol (The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the "[NC]" at the beginning of the line. For example, to generate a scene of falling leaves without any character, write: "[NC] The leaves are falling.")</h5>
+<h5 align="center">Tips: </h4>
+"""
+#################################################
+global attn_count, total_count, id_length, total_length, cur_step, cur_model_type
+global write
+global sa32, sa64
+global height, width
+attn_count = 0
+total_count = 0
+cur_step = 0
+id_length = 4
+total_length = 5
+cur_model_type = ""
+global attn_procs, unet
+attn_procs = {}
+###
+write = False
+###
+sa32 = 0.5
+sa64 = 0.5
+height = 768
+width = 768
+###
+global pipe
+global sd_model_path
+pipe = None
+sd_model_path = models_dict["Unstable"]["path"]  # "SG161222/RealVisXL_V4.0"
+single_files = models_dict["Unstable"]["single_files"]
+### LOAD Stable Diffusion Pipeline
+if single_files:
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        sd_model_path, torch_dtype=torch.float16
+    )
+else:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        sd_model_path, torch_dtype=torch.float16, use_safetensors=False
+    )
+print("pipE.device = ", device)
+pipe = pipe.to(device)
+pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler.set_timesteps(50)
+pipe.enable_vae_slicing()
+if device != "mps":
+    pipe.enable_model_cpu_offload()
+unet = pipe.unet
+cur_model_type = "Unstable" + "-" + "original"
+### Insert PairedAttention
+for name in unet.attn_processors.keys():
+    cross_attention_dim = (
+        None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+    )
+    if name.startswith("mid_block"):
+        hidden_size = unet.config.block_out_channels[-1]
+    elif name.startswith("up_blocks"):
+        block_id = int(name[len("up_blocks.")])
+        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+    elif name.startswith("down_blocks"):
+        block_id = int(name[len("down_blocks.")])
+        hidden_size = unet.config.block_out_channels[block_id]
+    if cross_attention_dim is None and (name.startswith("up_blocks")):
+        attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+        total_count += 1
+    else:
+        attn_procs[name] = AttnProcessor()
+print("successsfully load paired self-attention")
+print(f"number of the processor : {total_count}")
+unet.set_attn_processor(copy.deepcopy(attn_procs))
+global mask1024, mask4096
+mask1024, mask4096 = cal_attn_mask_xl(
+    total_length,
+    id_length,
+    sa32,
+    sa64,
+    height,
+    width,
+    device=device,
+    dtype=torch.float16,
+)
+
+######### Gradio Fuction #############
+
+
+def swap_to_gallery(images):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def upload_example_to_gallery(images, prompt, style, negative_prompt):
+    return (
+        gr.update(value=images, visible=True),
+        gr.update(visible=True),
+        gr.update(visible=False),
+    )
+
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+
+def remove_tips():
+    return gr.update(visible=False)
+
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive)
+
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [
+        p.replace("{prompt}", positive) for positive in positives
+    ], n + " " + negative
+
+
+def change_visiale_by_model_type(_model_type):
+    if _model_type == "Only Using Textual Description":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif _model_type == "Using Ref Images":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    else:
+        raise ValueError("Invalid model type", _model_type)
+
+
+def load_character_files(character_files: str):
+    if character_files == "":
+        raise gr.Error("Please set a character file!")
+    character_files_arr = character_files.splitlines()
+    primarytext = []
+    for character_file_name in character_files_arr:
+        character_file = torch.load(
+            character_file_name, map_location=torch.device("cpu")
+        )
+        primarytext.append(character_file["character"] + character_file["description"])
+    return array2string(primarytext)
+
+
+def load_character_files_on_running(unet, character_files: str):
+    if character_files == "":
+        return False
+    character_files_arr = character_files.splitlines()
+    for character_file in character_files_arr:
+        load_single_character_weights(unet, character_file)
+    return True
+
+
+######### Image Generation ##############
+def process_generation(
+    _sd_type,
+    _model_type,
+    _upload_images,
+    _num_steps,
+    style_name,
+    _Ip_Adapter_Strength,
+    _style_strength_ratio,
+    guidance_scale,
+    seed_,
+    sa32_,
+    sa64_,
+    id_length_,
+    general_prompt,
+    negative_prompt,
+    prompt_array,
+    G_height,
+    G_width,
+    _comic_type,
+    font_choice,
+    _char_files,
+):  # Corrected font_choice usage
+    if len(general_prompt.splitlines()) >= 3:
+        raise gr.Error(
+            "Support for more than three characters is temporarily unavailable due to VRAM limitations, but this issue will be resolved soon."
+        )
+    _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    if _model_type == "Photomaker" and "img" not in general_prompt:
+        raise gr.Error(
+            'Please add the triger word " img "  behind the class word you want to customize, such as: man img or woman img'
+        )
+    if _upload_images is None and _model_type != "original":
+        raise gr.Error(f"Cannot find any input face image!")
+    global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type
+    global write
+    global cur_step, attn_count
+    global height, width
+    height = G_height
+    width = G_width
+    global pipe
+    global sd_model_path, models_dict
+    sd_model_path = models_dict[_sd_type]
+    use_safe_tensor = True
+    for attn_processor in pipe.unet.attn_processors.values():
+        if isinstance(attn_processor, SpatialAttnProcessor2_0):
+            for values in attn_processor.id_bank.values():
+                del values
+            attn_processor.id_bank = {}
+            attn_processor.id_length = id_length
+            attn_processor.total_length = id_length + 1
+    gc.collect()
+    torch.cuda.empty_cache()
+    if cur_model_type != _sd_type + "-" + _model_type:
+        # apply the style template
+        ##### load pipe
+        del pipe
+        gc.collect()
+        if device == "cuda":
+            torch.cuda.empty_cache()
+        model_info = models_dict[_sd_type]
+        model_info["model_type"] = _model_type
+        print("device = ", device)
+        pipe = load_models(model_info, device=device, photomaker_path=photomaker_path)
+        set_attention_processor(pipe.unet, id_length_, is_ipadapter=False)
+        ##### ########################
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        cur_model_type = _sd_type + "-" + _model_type
+        pipe.enable_vae_slicing()
+        if device != "mps":
+            pipe.enable_model_cpu_offload()
+    else:
+        unet = pipe.unet
+        # unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+    load_chars = load_character_files_on_running(unet, character_files=_char_files)
+
+    prompts = prompt_array.splitlines()
+    global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals
+    character_dict, character_list = character_to_dict(general_prompt)
+
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    generator = torch.Generator(device=device).manual_seed(seed_)
+    sa32, sa64 = sa32_, sa64_
+    id_length = id_length_
+    clipped_prompts = prompts[:]
+    nc_indexs = []
+    for ind, prompt in enumerate(clipped_prompts):
+        if "[NC]" in prompt:
+            nc_indexs.append(ind)
+            if ind < id_length:
+                raise gr.Error(
+                    f"The first {id_length} row is id prompts, cannot use [NC]!"
+                )
+    prompts = [
+        prompt if "[NC]" not in prompt else prompt.replace("[NC]", "")
+        for prompt in clipped_prompts
+    ]
+
+    prompts = [
+        prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+    ]
+    print(prompts)
+    # id_prompts = prompts[:id_length]
+    (
+        character_index_dict,
+        invert_character_index_dict,
+        replace_prompts,
+        ref_indexs_dict,
+        ref_totals,
+    ) = process_original_prompt(character_dict, prompts.copy(), id_length)
+    if _model_type != "original":
+        input_id_images_dict = {}
+        if len(_upload_images) != len(character_dict.keys()):
+            raise gr.Error(
+                f"You upload images({len(_upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!"
+            )
+        for ind, img in enumerate(_upload_images):
+            input_id_images_dict[character_list[ind]] = [load_image(img)]
+    print(character_dict)
+    print(character_index_dict)
+    print(invert_character_index_dict)
+    # real_prompts = prompts[id_length:]
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+
+    attn_count = 0
+    # id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt)
+    # print(id_prompts)
+    setup_seed(seed_)
+    total_results = []
+    id_images = []
+    results_dict = {}
+    global cur_character
+    
+    if not load_chars:
+        for character_key in character_dict.keys():
+            cur_character = [character_key]
+            ref_indexs = ref_indexs_dict[character_key]
+            print(character_key, ref_indexs)
+            current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs]
+            print(current_prompts)
+            setup_seed(seed_)
+            generator = torch.Generator(device=device).manual_seed(seed_)
+            cur_step = 0
+            cur_positive_prompts, negative_prompt = apply_style(
+                style_name, current_prompts, negative_prompt
+            )
+            if _model_type == "original":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            elif _model_type == "Photomaker":
+                id_images = pipe(
+                    cur_positive_prompts,
+                    input_id_images=input_id_images_dict[character_key],
+                    num_inference_steps=_num_steps,
+                    guidance_scale=guidance_scale,
+                    start_merge_step=start_merge_step,
+                    height=height,
+                    width=width,
+                    negative_prompt=negative_prompt,
+                    generator=generator,
+                ).images
+            else:
+                raise NotImplementedError(
+                    "You should choice between original and Photomaker!",
+                    f"But you choice {_model_type}",
+                )
+
+            # total_results = id_images + total_results
+            # yield total_results
+            print(id_images)
+            for ind, img in enumerate(id_images):
+                print(ref_indexs[ind])
+                results_dict[ref_indexs[ind]] = img
+            # real_images = []
+            yield [results_dict[ind] for ind in results_dict.keys()]
+    write = False
+    if not load_chars:
+        real_prompts_inds = [
+            ind for ind in range(len(prompts)) if ind not in ref_totals
+        ]
+    else:
+        real_prompts_inds = [ind for ind in range(len(prompts))]
+    print(real_prompts_inds)
+
+    for real_prompts_ind in real_prompts_inds:
+        real_prompt = replace_prompts[real_prompts_ind]
+        cur_character = get_ref_character(prompts[real_prompts_ind], character_dict)
+        print(cur_character, real_prompt)
+        setup_seed(seed_)
+        if len(cur_character) > 1 and _model_type == "Photomaker":
+            raise gr.Error(
+                "Temporarily Not Support Multiple character in Ref Image Mode!"
+            )
+        generator = torch.Generator(device=device).manual_seed(seed_)
+        cur_step = 0
+        real_prompt = apply_style_positive(style_name, real_prompt)
+        if _model_type == "original":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images[0]
+        elif _model_type == "Photomaker":
+            results_dict[real_prompts_ind] = pipe(
+                real_prompt,
+                input_id_images=(
+                    input_id_images_dict[cur_character[0]]
+                    if real_prompts_ind not in nc_indexs
+                    else input_id_images_dict[character_list[0]]
+                ),
+                num_inference_steps=_num_steps,
+                guidance_scale=guidance_scale,
+                start_merge_step=start_merge_step,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+                nc_flag=True if real_prompts_ind in nc_indexs else False,
+            ).images[0]
+        else:
+            raise NotImplementedError(
+                "You should choice between original and Photomaker!",
+                f"But you choice {_model_type}",
+            )
+        yield [results_dict[ind] for ind in results_dict.keys()]
+    total_results = [results_dict[ind] for ind in range(len(prompts))]
+    if _comic_type != "No typesetting (default)":
+        captions = prompt_array.splitlines()
+        captions = [caption.replace("[NC]", "") for caption in captions]
+        captions = [
+            caption.split("#")[-1] if "#" in caption else caption
+            for caption in captions
+        ]
+        font_path = os.path.join("fonts", font_choice)
+        font = ImageFont.truetype(font_path, int(45))
+        total_results = (
+            get_comic(total_results, _comic_type, captions=captions, font=font)
+            + total_results
+        )
+    save_results(pipe.unet, total_results)
+
+    yield total_results
+
+
+def array2string(arr):
+    stringtmp = ""
+    for i, part in enumerate(arr):
+        if i != len(arr) - 1:
+            stringtmp += part + "\n"
+        else:
+            stringtmp += part
+
+    return stringtmp
+
+
+#################################################
+#################################################
+### define the interface
+
+css = """
+:root {
+    --main-blue: #4A90E2;
+    --tech-purple: #9B59B6;
+    --fresh-green: #2ECC71;
+    --bg-gradient: linear-gradient(135deg, #F5F7FA 0%, #E8F4FF 100%);
+}
+
+body {
+    background: var(--bg-gradient);
+    min-height: 100vh;
+}
+
+.gr-container {
+    max-width: 1200px!important;
+    margin: 0 auto!important;
+    gap: 40px!important;
+}
+
+.upload-section {
+    border: 2px dashed var(--main-blue)!important;
+    border-radius: 20px!important;
+    padding: 30px!important;
+}
+
+.generate-btn {
+    background: var(--fresh-green)!important;
+    color: white!important;
+    border-radius: 12px!important;
+}
+
+/* 其他样式保持与提供的一致 */
+"""
+
+# 提示词映射字典
+PROMPT_TEMPLATES = {
+    "双人冒险故事 🔥": {
+        "general": "[Bob] A man, wearing a black suit\n[Alice]a woman, wearing a white shirt",
+        "negative": "bad anatomy, bad hands, missing fingers...",  # 保持原负面提示词
+        "scenes": [
+            "[Bob] at home, read new paper #at home...",
+            "[Bob] on the road, near the forest",
+            # ... 其他场景
+        ]
+    },
+    "夜间森林探险 🌲": {
+        "general": "[Bob] A man img, wearing a black suit...",
+        "negative": "bad anatomy...",
+        "scenes": [...] 
+    },
+    # 其他模板...
+}
+
+def array2string(arr):
+    return "\n".join(arr)
+
+def load_example(prompt_key, style, files, height, width):
+    """处理示例加载"""
+    template = PROMPT_TEMPLATES[prompt_key]
+    return {
+        general_prompt: template["general"],
+        negative_prompt: template["negative"],
+        prompt_array: array2string(template["scenes"]),
+        style: style,
+        files: files,
+        G_height: height,
+        G_width: width
+    }
+
+with gr.Blocks(css=css, title="AI研学旅记") as demo:
+    # 状态变量
+    binary_matrixes = gr.State([])
+    color_layout = gr.State([])
+    
+    # 标题
+    gr.Markdown("<h1 class='main-title'>我的AI研学旅记</h1>")
+    
+    with gr.Row(elem_id="container"):
+        # 左侧输入区
+        with gr.Column(elem_classes="input-area"):
+            # === 文件上传 ===
+            with gr.Group(elem_classes="upload-section"):
+                files = gr.Files(
+                    label="上传研学照片",
+                    file_types=["image"],
+                    file_count="multiple",
+                    elem_id="fileInput"
+                )
+                uploaded_files = gr.Gallery(
+                    label="您的照片",
+                    columns=5,
+                    rows=1,
+                    height=200,
+                    visible=False
+                )
+            
+            # === 提示词选择 ===
+            prompt_btns = gr.Radio(
+                choices=list(PROMPT_TEMPLATES.keys()),
+                label="场景模板",
+                interactive=True,
+                elem_classes="prompt-section"
+            )
+            
+            # === 风格选择 ===
+            style = gr.Dropdown(
+                choices=["🎞️ ���本动漫风", "🌸 电影影视风", "🎨 摄影写真风", "🌟 漫画书风", "🌙 皮克斯/迪士尼角色", "📘 线条艺术风"],
+                value="🎞️ 日本动漫风",
+                label="艺术风格",
+                elem_classes="style-select"
+            )
+            
+            # === 固定参数 ===
+            with gr.Accordion("高级设置", open=False):
+                sa32_ = gr.Slider(visible=False, value=0.5)
+                sa64_ = gr.Slider(visible=False, value=0.5)
+                id_length_ = gr.Slider(visible=False, value=3)
+                guidance_scale = gr.Slider(visible=False, value=5)
+                num_steps = gr.Slider(visible=False, value=35)
+            
+            # === 生成按钮 ===
+            generate_btn = gr.Button("✨ 立即生成", elem_classes="generate-btn")
+
+        # 右侧输出区
+        with gr.Column(elem_classes="output-section"):
+            out_image = gr.Gallery(
+                label="生成结果",
+                columns=2,
+                height="auto",
+                elem_classes="output-image"
+            )
+            loading = gr.HTML("""
+                <div class="loading-overlay">
+                    <div style="text-align: center">
+                        <div style="font-size: 2rem">🎨 AI创作中...</div>
+                        <div style="font-size: 1.2rem; color: #666">正在将回忆转化为数字艺术</div>
+                    </div>
+                </div>
+            """, visible=False)
+
+    # === 事件处理 ===
+    files.upload(
+        fn=swap_to_gallery,
+        inputs=files,
+        outputs=[uploaded_files, files]
+    )
+
+    generate_btn.click(
+        fn=lambda: gr.update(visible=True),
+        outputs=loading
+    ).then(
+        fn=process_generation,
+        inputs=[
+            prompt_btns, style, files,
+            sa32_, sa64_, id_length_,
+            guidance_scale, num_steps,
+            gr.State(768), gr.State(768)  # 固定尺寸
+        ],
+        outputs=out_image
+    ).then(
+        fn=lambda: gr.update(visible=False),
+        outputs=loading
+    )
+
+    # === 示例 ===
+    gr.Examples(
+        examples=[
+            [
+                "双人冒险故事 🔥",
+                "🌟 漫画书风",
+                get_image_path_list("examples/taylor"),
+                768,
+                768
+            ],
+            [
+                "夜间森林探险 🌲", 
+                "🎞️ 日本动漫风",
+                get_image_path_list("examples/twoperson"),
+                1024,
+                1024
+            ],
+            # 其他示例...
+        ],
+        inputs=[prompt_btns, style, files, gr.State(768), gr.State(768)],
+        outputs=out_image,
+        fn=load_example,
+        label="预设场景"
+    )
+
+# 保留原有启动设置
+demo.launch(server_name="0.0.0.0", share=True)

main.py

@@ -0,0 +1,25 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from diffusers import StableDiffusionPipeline
+import torch
+from PIL import Image
+import io
+import base64
+
+app = FastAPI()
+
+# 🔧 加载模型（你可以替换成你本地修改后的 StoryDiffusion Pipeline）
+pipe = StableDiffusionPipeline.from_pretrained("你的模型路径").to("cuda")
+
+class PromptInput(BaseModel):
+    prompt: str
+
+@app.post("/generate")
+def generate_image(data: PromptInput):
+    image = pipe(data.prompt).images[0]
+
+    # 把图像编码为 base64，便于前端显示
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return {"image_base64": img_str}

page.html

@@ -0,0 +1,60 @@
+<!DOCTYPE html>
+<html lang="zh-CN">
+<head>
+    <meta charset="UTF-8">
+    <title>跳转页面</title>
+    <style>
+        @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@500&display=swap');
+
+        body {
+            margin: 0;
+            padding: 0;
+            height: 100vh;
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            background: linear-gradient(135deg, #0f0c29, #302b63, #24243e);
+            font-family: 'Orbitron', sans-serif;
+            color: #00ffe7;
+            overflow: hidden;
+        }
+
+        button {
+            font-size: 24px;
+            padding: 15px 30px;
+            background: rgba(0, 255, 231, 0.2);
+            border: 2px solid #00ffe7;
+            border-radius: 12px;
+            color: #00ffe7;
+            cursor: pointer;
+            box-shadow: 0 0 15px #00ffe7;
+            transition: all 0.3s ease;
+            backdrop-filter: blur(5px);
+        }
+
+        button:hover {
+            background: rgba(0, 255, 231, 0.4);
+            box-shadow: 0 0 25px #00ffe7;
+            transform: scale(1.05);
+        }
+
+        .glow {
+            position: absolute;
+            width: 100%;
+            height: 100%;
+            background: radial-gradient(circle, rgba(0, 255, 231, 0.1) 0%, transparent 70%);
+            animation: pulse 5s infinite;
+        }
+
+        @keyframes pulse {
+            0% { transform: scale(1); opacity: 0.5; }
+            50% { transform: scale(1.2); opacity: 0.2; }
+            100% { transform: scale(1); opacity: 0.5; }
+        }
+    </style>
+</head>
+<body>
+    <div class="glow"></div>
+    <button onclick="window.location.href='http://localhost:7861/'">gogogo</button>
+</body>
+</html>

predict.py

@@ -0,0 +1,781 @@
+# Prediction interface for Cog ⚙️
+# https://cog.run/python
+
+import os
+import copy
+import random
+import subprocess
+import numpy as np
+import time
+import torch
+import torch.nn.functional as F
+from PIL import ImageFont
+from cog import BasePredictor, Input, Path, BaseModel
+from diffusers import StableDiffusionXLPipeline, DDIMScheduler
+from diffusers.utils import load_image
+
+from utils import PhotoMakerStableDiffusionXLPipeline
+from utils.style_template import styles
+from utils.gradio_utils import (
+    AttnProcessor2_0 as AttnProcessor,
+)  # with torch2 installed
+from utils.gradio_utils import cal_attn_mask_xl
+from utils.utils import get_comic
+
+MODEL_URL = "https://weights.replicate.delivery/default/HVision_NKU/StoryDiffusion.tar"
+MODEL_CACHE = "model_weights"
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "Japanese Anime"
+
+global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet
+global sa32, sa64
+global write
+global height, width
+
+
+"""
+# load and upload the weights to replicate.delivery for faster booting on Replicate
+models_dict = {
+    "RealVision": "SG161222/RealVisXL_V4.0",
+    "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y",
+}
+# photomaker_path =  hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")
+photomaker_path = f"{MODEL_CACHE}/PhotoMaker/photomaker-v1.bin"
+
+pipe_unstable = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+    models_dict["Unstable"],
+    torch_dtype=torch.float16,
+    use_safetensors=False,
+)
+pipe_unstable.save_pretrained(f"{MODEL_CACHE}/Unstable/stablediffusionapi/sdxl-unstable-diffusers-y")
+
+pipe_realvision = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+    models_dict["RealVision"], torch_dtype=torch.float16, use_safetensors=True
+)
+pipe_realvision.save_pretrained(f"{MODEL_CACHE}/RealVision/SG161222/RealVisXL_V4.0")
+"""
+
+
+class ModelOutput(BaseModel):
+    comic: Path
+    individual_images: list[Path]
+
+
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-x", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def apply_style_positive(style_name: str, positive: str):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive)
+
+
+def apply_style(style_name: str, positives: list, negative: str = ""):
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return [
+        p.replace("{prompt}", positive) for positive in positives
+    ], n + " " + negative
+
+
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global total_count
+    total_count = 0
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = (
+            None
+            if name.endswith("attn1.processor")
+            else unet.config.cross_attention_dim
+        )
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+                total_count += 1
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+    print("Successfully load paired self-attention")
+    print(f"Number of the processor : {total_count}")
+
+
+#################################################
+########Consistent Self-Attention################
+#################################################
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+################################################################################################################################################################################
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        id_length=4,
+        device="cuda",
+        dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        global total_count, attn_count, cur_step, mask1024, mask4096
+        global sa32, sa64
+        global write
+        global height, width
+        if write:
+            self.id_bank[cur_step] = [
+                hidden_states[: self.id_length],
+                hidden_states[self.id_length :],
+            ]
+        else:
+            encoder_hidden_states = torch.cat(
+                (
+                    self.id_bank[cur_step][0].to(self.device),
+                    hidden_states[:1],
+                    self.id_bank[cur_step][1].to(self.device),
+                    hidden_states[1:],
+                )
+            )
+        # skip in early step
+        if cur_step < 5:
+            hidden_states = self.__call2__(
+                attn, hidden_states, encoder_hidden_states, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            if random_number > rand_num:
+                if not write:
+                    if hidden_states.shape[1] == (height // 32) * (width // 32):
+                        attention_mask = mask1024[
+                            mask1024.shape[0] // self.total_length * self.id_length :
+                        ]
+                    else:
+                        attention_mask = mask4096[
+                            mask4096.shape[0] // self.total_length * self.id_length :
+                        ]
+                else:
+                    if hidden_states.shape[1] == (height // 32) * (width // 32):
+                        attention_mask = mask1024[
+                            : mask1024.shape[0] // self.total_length * self.id_length,
+                            : mask1024.shape[0] // self.total_length * self.id_length,
+                        ]
+                    else:
+                        attention_mask = mask4096[
+                            : mask4096.shape[0] // self.total_length * self.id_length,
+                            : mask4096.shape[0] // self.total_length * self.id_length,
+                        ]
+                hidden_states = self.__call1__(
+                    attn, hidden_states, encoder_hidden_states, attention_mask, temb
+                )
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            mask1024, mask4096 = cal_attn_mask_xl(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height,
+                width,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call1__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            total_batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                total_batch_size, channel, height * width
+            ).transpose(1, 2)
+        total_batch_size, nums_token, channel = hidden_states.shape
+        img_nums = total_batch_size // 2
+        hidden_states = hidden_states.view(-1, img_nums, nums_token, channel).reshape(
+            -1, img_nums * nums_token, channel
+        )
+
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(
+                -1, self.id_length + 1, nums_token, channel
+            ).reshape(-1, (self.id_length + 1) * nums_token, channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            total_batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                total_batch_size, channel, height, width
+            )
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        hidden_states = hidden_states / attn.rescale_output_factor
+        # print(hidden_states.shape)
+        return hidden_states
+
+    def __call2__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+        # print(hidden_states.shape)
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        else:
+            encoder_hidden_states = encoder_hidden_states.view(
+                -1, self.id_length + 1, sequence_length, channel
+            ).reshape(-1, (self.id_length + 1) * sequence_length, channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, 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)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+
+        models_dict = {
+            "RealVision": "SG161222/RealVisXL_V4.0",
+            "Unstable": "stablediffusionapi/sdxl-unstable-diffusers-y",
+        }
+
+        if not os.path.exists(MODEL_CACHE):
+            download_weights(MODEL_URL, MODEL_CACHE)
+
+        photomaker_path = f"{MODEL_CACHE}/PhotoMaker/photomaker-v1.bin"
+
+        self.sdxl_pipe_unstable = StableDiffusionXLPipeline.from_pretrained(
+            f"{MODEL_CACHE}/Unstable/sdxl/stablediffusionapi/sdxl-unstable-diffusers-y",
+            torch_dtype=torch.float16,
+        )
+        self.sdxl_pipe_realvision = StableDiffusionXLPipeline.from_pretrained(
+            f"{MODEL_CACHE}/RealVision/sdxl/SG161222/RealVisXL_V4.0",
+            torch_dtype=torch.float16,
+        )
+
+        self.pipe_unstable = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+            f"{MODEL_CACHE}/Unstable/stablediffusionapi/sdxl-unstable-diffusers-y",
+            torch_dtype=torch.float16,
+            use_safetensors=False,
+        )
+        self.pipe_unstable.load_photomaker_adapter(
+            os.path.dirname(photomaker_path),
+            subfolder="",
+            weight_name=os.path.basename(photomaker_path),
+            trigger_word="img",  # define the trigger word
+        )
+
+        self.pipe_realvision = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+            f"{MODEL_CACHE}/RealVision/SG161222/RealVisXL_V4.0",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        self.pipe_realvision.load_photomaker_adapter(
+            os.path.dirname(photomaker_path),
+            subfolder="",
+            weight_name=os.path.basename(photomaker_path),
+            trigger_word="img",  # define the trigger word
+        )
+        self.pipe_realvision.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        self.pipe_realvision.fuse_lora()
+
+    @torch.inference_mode()
+    def predict(
+        self,
+        sd_model: str = Input(
+            description="Choose a model",
+            choices=["Unstable", "RealVision"],
+            default="Unstable",
+        ),
+        ref_image: Path = Input(
+            description="Reference image for the character",
+            default=None,
+        ),
+        character_description: str = Input(
+            description="General description of the character. If ref_image above is provided, making sure to follow the class word you want to customize with the trigger word 'img', such as: 'man img' or 'woman img' or 'girl img'",
+            default="a man, wearing black suit",
+        ),
+        negative_prompt: str = Input(
+            description="Describe things you do not want to see in the output",
+            default="bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, animate, amputation, disconnected limbs",
+        ),
+        comic_description: str = Input(
+            description="Comic Description. Each frame is divided by a new line. Only the first 10 prompts are valid for demo speed! For comic_description NOT using ref_image: (1) Support Typesetting Style and Captioning. By default, the prompt is used as the caption for each image. If you need to change the caption, add a '#' at the end of each line. Only the part after the '#' will be added as a caption to the image. (2) The [NC] symbol is used as a flag to indicate that no characters should be present in the generated scene images. If you want do that, prepend the '[NC]' at the beginning of the line.",
+            default="at home, read new paper #at home, The newspaper says there is a treasure house in the forest.\non the road, near the forest\n[NC] The car on the road, near the forest #He drives to the forest in search of treasure.\n[NC]A tiger appeared in the forest, at night \nvery frightened, open mouth, in the forest, at night\nrunning very fast, in the forest, at night\n[NC] A house in the forest, at night #Suddenly, he discovers the treasure house!\nin the house filled with  treasure, laughing, at night #He is overjoyed inside the house.",
+        ),
+        style_name: str = Input(
+            description="Style template",
+            choices=STYLE_NAMES,
+            default=DEFAULT_STYLE_NAME,
+        ),
+        comic_style: str = Input(
+            description="Select the comic style for the combined comic",
+            choices=["Four Pannel", "Classic Comic Style"],
+            default="Classic Comic Style",
+        ),
+        style_strength_ratio: int = Input(
+            description="Style strength of Ref Image (%), only used if ref_image is provided",
+            default=20,
+            ge=15,
+            le=50,
+        ),
+        image_width: int = Input(
+            description="Width of output image",
+            choices=[
+                256,
+                288,
+                320,
+                352,
+                384,
+                416,
+                448,
+                480,
+                512,
+                544,
+                576,
+                608,
+                640,
+                672,
+                704,
+                736,
+                768,
+                800,
+                832,
+                864,
+                896,
+                928,
+                960,
+                992,
+                1024,
+            ],
+            default=768,
+        ),
+        image_height: int = Input(
+            description="Height of output image",
+            choices=[
+                256,
+                288,
+                320,
+                352,
+                384,
+                416,
+                448,
+                480,
+                512,
+                544,
+                576,
+                608,
+                640,
+                672,
+                704,
+                736,
+                768,
+                800,
+                832,
+                864,
+                896,
+                928,
+                960,
+                992,
+                1024,
+            ],
+            default=768,
+        ),
+        num_steps: int = Input(
+            description="Number of sample steps", ge=20, le=50, default=25
+        ),
+        guidance_scale: float = Input(
+            description="Scale for classifier-free guidance", ge=0.1, le=10, default=5
+        ),
+        seed: int = Input(
+            description="Random seed. Leave blank to randomize the seed", default=None
+        ),
+        sa32_setting: float = Input(
+            description="The degree of Paired Attention at 32 x 32 self-attention layers",
+            default=0.5,
+            ge=0,
+            le=1.0,
+        ),
+        sa64_setting: float = Input(
+            description="The degree of Paired Attention at 64 x 64 self-attention layers",
+            default=0.5,
+            ge=0,
+            le=1.0,
+        ),
+        num_ids: int = Input(
+            description="Number of id images in total images. This should not exceed total number of line-separated prompts",
+            default=3,
+        ),
+        output_format: str = Input(
+            description="Format of the output images",
+            choices=["webp", "jpg", "png"],
+            default="webp",
+        ),
+        output_quality: int = Input(
+            description="Quality of the output images, from 0 to 100. 100 is best quality, 0 is lowest quality",
+            default=80,
+            ge=0,
+            le=100,
+        ),
+    ) -> ModelOutput:
+        """Run a single prediction on the model"""
+
+        global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet
+        global sa32, sa64
+        global write
+        global height, width
+
+        assert (
+            len(character_description.strip()) > 0
+        ), "Please provide the description of the character."
+
+        if ref_image is not None:
+            assert (
+                "img" in character_description
+            ), f"When using ref_image, please add the trigger word 'img' behind the class word you want to customize, such as: man img or woman img"
+            assert (
+                "[NC]" not in comic_description
+            ), "You should not use trigger word [NC] when ref_image is provided."
+
+        height = image_height
+        width = image_width
+        id_length = num_ids
+        sa32 = sa32_setting
+        sa64 = sa64_setting
+
+        clipped_prompts = comic_description.splitlines()[:10]
+        print(clipped_prompts)
+        prompts = [
+            (
+                character_description + "," + prompt
+                if "[NC]" not in prompt
+                else prompt.replace("[NC]", "")
+            )
+            for prompt in clipped_prompts
+        ]
+        print(prompts)
+        prompts = [
+            prompt.rpartition("#")[0].strip() if "#" in prompt else prompt.strip()
+            for prompt in prompts
+        ]
+        print(prompts)
+        assert id_length <= len(
+            prompts
+        ), "id_length should not exceed total number of line-separated prompts"
+
+        id_prompts = prompts[:id_length]
+        real_prompts = prompts[id_length:]
+
+        if seed is None:
+            seed = int.from_bytes(os.urandom(2), "big")
+        print(f"Using seed: {seed}")
+################################################################################################################################################################################
+        device = "cuda"
+        setup_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
+
+        torch.cuda.empty_cache()
+
+        model_type = "original" if ref_image is None else "Photomaker"
+
+        if model_type == "original":
+            pipe = (
+                self.sdxl_pipe_realvision
+                if style_name == "(No style)"
+                else self.sdxl_pipe_unstable
+            )
+            pipe = pipe.to(device)
+            pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        else:
+            if sd_model != "RealVision" and style_name != "(No style)":
+                pipe = self.pipe_unstable.to(device)
+            else:
+                pipe = self.pipe_realvision.to(device)
+            pipe.id_encoder.to(device)
+
+        write = True
+        cur_step = 0
+        attn_count = 0
+
+        set_attention_processor(pipe.unet, id_length, is_ipadapter=False)
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+        curmodel_type = sd_model + "-" + model_type + "" + str(id_length)
+
+        id_prompts, negative_prompt = apply_style(
+            style_name, id_prompts, negative_prompt
+        )
+
+        total_results = []
+        if model_type == "original":
+            id_images = pipe(
+                id_prompts,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images
+        else:
+            input_id_images = [load_image(str(ref_image))]
+            start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)
+            id_images = pipe(
+                id_prompts,
+                input_id_images=input_id_images,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                start_merge_step=start_merge_step,
+                height=height,
+                width=width,
+                negative_prompt=negative_prompt,
+                generator=generator,
+            ).images
+
+        total_results = id_images + total_results
+
+        real_images = []
+        write = False
+        for real_prompt in real_prompts:
+            cur_step = 0
+            real_prompt = apply_style_positive(style_name, real_prompt)
+            if model_type == "original":
+                real_images.append(
+                    pipe(
+                        real_prompt,
+                        num_inference_steps=num_steps,
+                        guidance_scale=guidance_scale,
+                        height=height,
+                        width=width,
+                        negative_prompt=negative_prompt,
+                        generator=generator,
+                    ).images[0]
+                )
+            else:
+                real_images.append(
+                    pipe(
+                        real_prompt,
+                        input_id_images=input_id_images,
+                        num_inference_steps=num_steps,
+                        guidance_scale=guidance_scale,
+                        start_merge_step=start_merge_step,
+                        height=height,
+                        width=width,
+                        negative_prompt=negative_prompt,
+                        generator=generator,
+                    ).images[0]
+                )
+
+            total_results = [real_images[-1]] + total_results
+
+        captions = clipped_prompts
+        captions = [caption.replace("[NC]", "") for caption in captions]
+        captions = [
+            caption.split("#")[-1].strip() if "#" in caption else caption.strip()
+            for caption in captions
+        ]
+
+        comic = get_comic(
+            id_images + real_images,
+            comic_style,
+            captions=captions,
+            font=ImageFont.truetype("./fonts/Inkfree.ttf", int(45)),
+        )
+
+        extension = output_format.lower()
+        extension = "jpeg" if extension == "jpg" else extension
+        comic_out = f"/tmp/comic.{extension}"
+        comic[0].save(comic_out)
+
+        save_params = {"format": extension.upper()}
+        if not output_format == "png":
+            save_params["quality"] = output_quality
+            save_params["optimize"] = True
+
+        output_paths = []
+        for index, sample in enumerate(total_results[::-1]):
+            output_filename = f"/tmp/out-{index}.{extension}"
+            sample.save(output_filename, **save_params)
+            output_paths.append(Path(output_filename))
+
+        del pipe
+
+        return ModelOutput(comic=Path(comic_out), individual_images=output_paths)

requirements.txt

@@ -0,0 +1,15 @@
+gradio==4.22.0
+xformers==0.0.20
+torch==2.0.1
+torchvision==0.15.2
+diffusers==0.25.0
+transformers==4.36.2
+huggingface-hub==0.20.2
+spaces==0.19.4
+numpy
+accelerate
+safetensors
+omegaconf
+peft
+httpx==0.27.0
+safetensors==0.4.0

run.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+echo "✨ 启动 Gradio（本地监听 7860 端口）"
+python gradio_app_sdxl_specific_id_low_vram.py &  # 后台运行
+
+sleep 60
+
+echo "🌐 尝试使用 cloudflared 暴露端口..."
+cloudflared tunnel --url http://localhost:7860

test.py

@@ -0,0 +1,4 @@
+import torch
+print("CUDA available:", torch.cuda.is_available())
+print("cuDNN version:", torch.backends.cudnn.version())
+print("cuDNN enabled:", torch.backends.cudnn.enabled)

update.md

@@ -0,0 +1,28 @@
+## Update History
+
+### Update 2023-05-14
+
+- Support Two persons,support for more characters will also be possible in the feature. In Pnhotomaker, currently, only one person can appear in a single image.
+- Auto Save generated images in the ‘results’ folder.
+- I have changed the way to fill in prompts; please refer to the example provided.
+
+### Update 2024-05-08
+
+- Support [NC] in Ref Image Model  (Photomaker work best in 1024x1024 but may cost a lot of GPU memory, I recommend you to use the res. as larger as possible)
+
+<img src="results_examples/image1.png" height=100>
+
+- Merge Push by @cryptowooser to support lastest pillow. But you may be updated pillow if you using the old version.
+
+
+
+### Todo 
+
+- Support add captions on all images for the classical commic Typesetting Style
+
+
+
+
+### Welcome to contribute
+
+- Various layout styles.