Final Files

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

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.

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.py

@@ -0,0 +1,800 @@
+from email.policy import default
+from this import d
+import gradio as gr
+import numpy as np
+import torch
+import gc
+from huggingface_hub import hf_hub_download
+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
+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,cal_attn_indice_xl_effcient_memory
+import copy
+import os
+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
+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, indices1024,indices4096
+        global sa32, sa64
+        global write
+        global height,width
+        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:
+            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)
+            self.id_bank[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)))
+            encoder_arr = [tensor.to(self.device)  for tensor in self.id_bank[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):
+                        encoder_hidden_states_tmp = torch.cat(encoder_arr[0:img_ind] + encoder_arr[img_ind+1:] + [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)
+                    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 __call1__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        attn_indices = 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 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>
+'''
+
+
+#################################################
+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>[email protected]</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.")</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 = ""
+device="cuda"
+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"]#"SG161222/RealVisXL_V4.0"
+### LOAD Stable Diffusion Pipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16, use_safetensors = False)
+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()
+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)
+
+
+######### 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):
+    _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()
+    if cur_model_type != _sd_type+"-"+_model_type:
+        if _sd_type == "Unstable":
+            use_safe_tensor = False
+        # apply the style template
+        ##### load pipe
+        del pipe
+        gc.collect()
+        torch.cuda.empty_cache()
+        if _model_type == "original":
+            pipe = StableDiffusionXLPipeline.from_pretrained(sd_model_path, torch_dtype=torch.float16, use_safetensors=use_safe_tensor)
+            pipe = pipe.to(device)
+            set_attention_processor(pipe.unet,id_length_,is_ipadapter = False)
+        elif _model_type == "Photomaker":
+            pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+                sd_model_path, torch_dtype=torch.float16, use_safetensors=use_safe_tensor)
+            pipe = pipe.to(device)
+            pipe.load_photomaker_adapter(
+                os.path.dirname(photomaker_path),
+                subfolder="",
+                weight_name=os.path.basename(photomaker_path),
+                trigger_word="img"  # define the trigger word
+            )
+            pipe.fuse_lora()
+            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
+        pipe.enable_vae_slicing()
+        pipe.enable_model_cpu_offload()
+    else:
+        unet = pipe.unet
+        # unet.set_attn_processor(copy.deepcopy(attn_procs))
+    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
+    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='', 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):
+                    sa32_ = gr.Slider(label=" (The degree of Paired Attention at 32 x 32 self-attention layers) ", minimum=0, maximum=1., 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., value=0.5, step=0.1)
+                    id_length_ = gr.Slider(label= "Number of id images in total images" , minimum=2, maximum=4, value=2, 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=[
+            [0,0.5,0.5,2,"a man, wearing black suit",
+                "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, read new paper #at home, The newspaper says there is a treasure house in the forest.",
+                            "on the road, near the forest",
+                            "[NC] The car on the road, near the forest #He drives to the forest in search of treasure.",
+                            "[NC]A tiger appeared in the forest, at night ",
+                            "very frightened, open mouth, in the forest, at night",
+                            "running very fast, in the forest, at night",
+                            "[NC] A house in the forest, at night #Suddenly, he discovers the treasure house!",
+                            "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
+            ],
+            [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,3,"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,3,"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.",
+                            "[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"]),
+                    "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 = False)

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 character_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:0"
+        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)