Support PuLID (#2838)

* Add preprocessors

* Fix resolution param

* Fix various issues

* Add PuLID attn

* remove unused import

* Resize img before passing to facexlib

* safe unload

internal_controlnet/external_code.py

@@ -11,7 +11,7 @@
 from modules.safe import unsafe_torch_load
 from scripts import global_state
 from scripts.logging import logger
-from scripts.enums import HiResFixOption
+from scripts.enums import HiResFixOption, PuLIDMode
 from scripts.supported_preprocessor import Preprocessor, PreprocessorParameter
 
 from modules.api import api
@@ -207,6 +207,10 @@ class ControlNetUnit:
     # The effective region mask that unit's effect should be restricted to.
     effective_region_mask: Optional[np.ndarray] = None
 
+    # The weight mode for PuLID.
+    # https://github.com/ToTheBeginning/PuLID
+    pulid_mode: PuLIDMode = PuLIDMode.FIDELITY
+
     # The tensor input for ipadapter. When this field is set in the API,
     # the base64string will be interpret by torch.load to reconstruct ipadapter
     # preprocessor output.
@@ -243,6 +247,7 @@ def infotext_excluded_fields() -> List[str]:
             # provide much information when restoring the unit.
             "inpaint_crop_input_image",
             "effective_region_mask",
+            "pulid_mode",
         ]
 
     @property

requirements.txt

@@ -6,3 +6,4 @@ addict
 yapf
 albumentations==1.4.3
 matplotlib
+facexlib

scripts/api.py

@@ -179,7 +179,7 @@ def accept(self, json_dict: dict) -> None:
                 low_vram=low_vram,
             )
             if preprocessor.returns_image:
-                images.append(encode_to_base64(result.display_image))
+                images.append(encode_to_base64(result.display_images[0]))
             else:
                 tensors.append(encode_tensor_to_base64(result.value))
 

scripts/controlnet.py

@@ -17,12 +17,14 @@
 import scripts.preprocessor as preprocessor_init  # noqa
 from annotator.util import HWC3
 from scripts import global_state, hook, external_code, batch_hijack, controlnet_version, utils
+from internal_controlnet.external_code import ControlMode
 from scripts.controlnet_lora import bind_control_lora, unbind_control_lora
 from scripts.controlnet_lllite import clear_all_lllite
 from scripts.ipadapter.plugable_ipadapter import ImageEmbed, clear_all_ip_adapter
+from scripts.ipadapter.pulid_attn import PULID_SETTING_FIDELITY, PULID_SETTING_STYLE
 from scripts.utils import load_state_dict, get_unique_axis0, align_dim_latent
 from scripts.hook import ControlParams, UnetHook, HackedImageRNG
-from scripts.enums import ControlModelType, StableDiffusionVersion, HiResFixOption
+from scripts.enums import ControlModelType, StableDiffusionVersion, HiResFixOption, PuLIDMode
 from scripts.controlnet_ui.controlnet_ui_group import ControlNetUiGroup, UiControlNetUnit
 from scripts.controlnet_ui.photopea import Photopea
 from scripts.logging import logger
@@ -279,6 +281,7 @@ def preprocess_input_image(input_image: np.ndarray):
         )
         detected_map = result.value
         is_image = preprocessor.returns_image
+        # TODO: Refactor img control detection logic.
         if high_res_fix:
             if is_image:
                 hr_control, hr_detected_map = Script.detectmap_proc(detected_map, unit.module, resize_mode, hr_y, hr_x)
@@ -293,7 +296,8 @@ def preprocess_input_image(input_image: np.ndarray):
             store_detected_map(detected_map, unit.module)
         else:
             control = detected_map
-            store_detected_map(input_image, unit.module)
+            for image in result.display_images:
+                store_detected_map(image, unit.module)
 
         if control_model_type == ControlModelType.T2I_StyleAdapter:
             control = control['last_hidden_state']
@@ -1092,8 +1096,8 @@ def ad_process_control(cc: List[torch.Tensor], cn_ad_keyframe_idx=cn_ad_keyframe
                 global_average_pooling=global_average_pooling,
                 hr_hint_cond=hr_control,
                 hr_option=HiResFixOption.from_value(unit.hr_option) if high_res_fix else HiResFixOption.BOTH,
-                soft_injection=control_mode != external_code.ControlMode.BALANCED,
-                cfg_injection=control_mode == external_code.ControlMode.CONTROL,
+                soft_injection=control_mode != ControlMode.BALANCED,
+                cfg_injection=control_mode == ControlMode.CONTROL,
                 effective_region_mask=(
                     get_pytorch_control(unit.effective_region_mask)[:, 0:1, :, :]
                     if unit.effective_region_mask is not None
@@ -1190,7 +1194,7 @@ def recolor_intensity_post_processing(x, i):
 
         is_low_vram = any(unit.low_vram for unit in self.enabled_units)
 
-        for i, param in enumerate(forward_params):
+        for i, (param, unit) in enumerate(zip(forward_params, self.enabled_units)):
             if param.control_model_type == ControlModelType.IPAdapter:
                 if param.advanced_weighting is not None:
                     logger.info(f"IP-Adapter using advanced weighting {param.advanced_weighting}")
@@ -1205,6 +1209,13 @@ def recolor_intensity_post_processing(x, i):
                     weight = param.weight
 
                 h, w, hr_y, hr_x = Script.get_target_dimensions(p)
+                pulid_mode = PuLIDMode(unit.pulid_mode)
+                if pulid_mode == PuLIDMode.STYLE:
+                    pulid_attn_setting = PULID_SETTING_STYLE
+                else:
+                    assert pulid_mode == PuLIDMode.FIDELITY
+                    pulid_attn_setting = PULID_SETTING_FIDELITY
+
                 param.control_model.hook(
                     model=unet,
                     preprocessor_outputs=param.hint_cond,
@@ -1215,6 +1226,7 @@ def recolor_intensity_post_processing(x, i):
                     latent_width=w // 8,
                     latent_height=h // 8,
                     effective_region_mask=param.effective_region_mask,
+                    pulid_attn_setting=pulid_attn_setting,
                 )
             if param.control_model_type == ControlModelType.Controlllite:
                 param.control_model.hook(

scripts/controlnet_ui/controlnet_ui_group.py

@@ -18,7 +18,7 @@
 from scripts.controlnet_ui.preset import ControlNetPresetUI
 from scripts.controlnet_ui.photopea import Photopea
 from scripts.controlnet_ui.advanced_weight_control import AdvancedWeightControl
-from scripts.enums import InputMode
+from scripts.enums import InputMode, PuLIDMode
 from modules import shared
 from modules.ui_components import FormRow, FormHTML, ToolButton
 
@@ -287,6 +287,7 @@ def __init__(
         self.batch_image_dir_state = None
         self.output_dir_state = None
         self.advanced_weighting = gr.State(None)
+        self.pulid_mode = None
 
         # API-only fields
         self.ipadapter_input = gr.State(None)
@@ -626,6 +627,15 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             visible=False,
         )
 
+        self.pulid_mode = gr.Radio(
+            choices=[e.value for e in PuLIDMode],
+            value=self.default_unit.pulid_mode.value,
+            label="PuLID Mode",
+            elem_id=f"{elem_id_tabname}_{tabname}_controlnet_pulid_mode_radio",
+            elem_classes="controlnet_pulid_mode_radio",
+            visible=False,
+        )
+
         self.loopback = gr.Checkbox(
             label="[Batch Loopback] Automatically send generated images to this ControlNet unit in batch generation",
             value=self.default_unit.loopback,
@@ -673,6 +683,7 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             self.save_detected_map,
             self.advanced_weighting,
             self.effective_region_mask,
+            self.pulid_mode,
         )
 
         unit = gr.State(self.default_unit)
@@ -947,7 +958,7 @@ def is_openpose(module: str):
 
             return (
                 # Update to `generated_image`
-                gr.update(value=result.display_image, visible=True, interactive=False),
+                gr.update(value=result.display_images[0], visible=True, interactive=False),
                 # preprocessor_preview
                 gr.update(value=True),
                 # openpose editor
@@ -1118,6 +1129,14 @@ def register_shift_upload_mask(self):
             show_progress=False,
         )
 
+    def register_shift_pulid_mode(self):
+        self.model.change(
+            fn=lambda model: gr.update(visible="pulid" in model.lower()),
+            inputs=[self.model],
+            outputs=[self.pulid_mode],
+            show_progress=False,
+        )
+
     def register_sync_batch_dir(self):
         def determine_batch_dir(batch_dir, fallback_dir, fallback_fallback_dir):
             if batch_dir:
@@ -1220,6 +1239,7 @@ def register_core_callbacks(self):
         self.register_build_sliders()
         self.register_shift_preview()
         self.register_shift_upload_mask()
+        self.register_shift_pulid_mode()
         self.register_create_canvas()
         self.register_clear_preview()
         self.register_multi_images_upload()

scripts/enums.py

@@ -247,3 +247,8 @@ class InputMode(Enum):
     # Input is a directory. 1 generation. Each generation takes N input image
     # from the directory.
     MERGE = "merge"
+
+
+class PuLIDMode(Enum):
+    FIDELITY = "Fidelity"
+    STYLE = "Extremely style"

scripts/ipadapter/image_proj_models.py

@@ -269,3 +269,65 @@ def forward(self, x):
 
         latents = self.proj_out(latents)
         return self.norm_out(latents)
+
+
+class PuLIDEncoder(nn.Module):
+    def __init__(self, width=1280, context_dim=2048, num_token=5):
+        super().__init__()
+        self.num_token = num_token
+        self.context_dim = context_dim
+        h1 = min((context_dim * num_token) // 4, 1024)
+        h2 = min((context_dim * num_token) // 2, 1024)
+        self.body = nn.Sequential(
+            nn.Linear(width, h1),
+            nn.LayerNorm(h1),
+            nn.LeakyReLU(),
+            nn.Linear(h1, h2),
+            nn.LayerNorm(h2),
+            nn.LeakyReLU(),
+            nn.Linear(h2, context_dim * num_token),
+        )
+
+        for i in range(5):
+            setattr(
+                self,
+                f"mapping_{i}",
+                nn.Sequential(
+                    nn.Linear(1024, 1024),
+                    nn.LayerNorm(1024),
+                    nn.LeakyReLU(),
+                    nn.Linear(1024, 1024),
+                    nn.LayerNorm(1024),
+                    nn.LeakyReLU(),
+                    nn.Linear(1024, context_dim),
+                ),
+            )
+
+            setattr(
+                self,
+                f"mapping_patch_{i}",
+                nn.Sequential(
+                    nn.Linear(1024, 1024),
+                    nn.LayerNorm(1024),
+                    nn.LeakyReLU(),
+                    nn.Linear(1024, 1024),
+                    nn.LayerNorm(1024),
+                    nn.LeakyReLU(),
+                    nn.Linear(1024, context_dim),
+                ),
+            )
+
+    def forward(self, x, y):
+        # x shape [N, C]
+        x = self.body(x)
+        x = x.reshape(-1, self.num_token, self.context_dim)
+
+        hidden_states = ()
+        for i, emb in enumerate(y):
+            hidden_state = getattr(self, f"mapping_{i}")(emb[:, :1]) + getattr(
+                self, f"mapping_patch_{i}"
+            )(emb[:, 1:]).mean(dim=1, keepdim=True)
+            hidden_states += (hidden_state,)
+        hidden_states = torch.cat(hidden_states, dim=1)
+
+        return torch.cat([x, hidden_states], dim=1)

scripts/ipadapter/ipadapter_model.py

@@ -12,6 +12,7 @@
     MLPProjModel,
     MLPProjModelFaceId,
     ProjModelFaceIdPlus,
+    PuLIDEncoder,
 )
 
 
@@ -71,6 +72,7 @@ def __init__(
         is_faceid: bool,
         is_portrait: bool,
         is_instantid: bool,
+        is_pulid: bool,
         is_v2: bool,
     ):
         super().__init__()
@@ -85,9 +87,12 @@ def __init__(
         self.is_v2 = is_v2
         self.is_faceid = is_faceid
         self.is_instantid = is_instantid
+        self.is_pulid = is_pulid
         self.clip_extra_context_tokens = 16 if (self.is_plus or is_portrait) else 4
 
-        if is_instantid:
+        if self.is_pulid:
+            self.image_proj_model = PuLIDEncoder()
+        elif self.is_instantid:
             self.image_proj_model = self.init_proj_instantid()
         elif is_faceid:
             self.image_proj_model = self.init_proj_faceid()
@@ -235,6 +240,34 @@ def _get_image_embeds_instantid(
             self.image_proj_model(torch.zeros_like(prompt_image_emb)),
         )
 
+    def _get_image_embeds_pulid(self, pulid_proj_input) -> ImageEmbed:
+        """Get image embeds for pulid."""
+        id_cond = torch.cat(
+            [
+                pulid_proj_input.id_ante_embedding.to(
+                    device=self.device, dtype=torch.float32
+                ),
+                pulid_proj_input.id_cond_vit.to(
+                    device=self.device, dtype=torch.float32
+                ),
+            ],
+            dim=-1,
+        )
+        id_vit_hidden = [
+            t.to(device=self.device, dtype=torch.float32)
+            for t in pulid_proj_input.id_vit_hidden
+        ]
+        return ImageEmbed(
+            self.image_proj_model(
+                id_cond,
+                id_vit_hidden,
+            ),
+            self.image_proj_model(
+                torch.zeros_like(id_cond),
+                [torch.zeros_like(t) for t in id_vit_hidden],
+            ),
+        )
+
     @staticmethod
     def load(state_dict: dict, model_name: str) -> IPAdapterModel:
         """
@@ -245,6 +278,7 @@ def load(state_dict: dict, model_name: str) -> IPAdapterModel:
         is_v2 = "v2" in model_name
         is_faceid = "faceid" in model_name
         is_instantid = "instant_id" in model_name
+        is_pulid = "pulid" in model_name.lower()
         is_portrait = "portrait" in model_name
         is_full = "proj.3.weight" in state_dict["image_proj"]
         is_plus = (
@@ -256,8 +290,8 @@ def load(state_dict: dict, model_name: str) -> IPAdapterModel:
         sdxl = cross_attention_dim == 2048
         sdxl_plus = sdxl and is_plus
 
-        if is_instantid:
-            # InstantID does not use clip embedding.
+        if is_instantid or is_pulid:
+            # InstantID/PuLID does not use clip embedding.
             clip_embeddings_dim = None
         elif is_faceid:
             if is_plus:
@@ -291,10 +325,13 @@ def load(state_dict: dict, model_name: str) -> IPAdapterModel:
             is_portrait=is_portrait,
             is_instantid=is_instantid,
             is_v2=is_v2,
+            is_pulid=is_pulid,
         )
 
     def get_image_emb(self, preprocessor_output) -> ImageEmbed:
-        if self.is_instantid:
+        if self.is_pulid:
+            return self._get_image_embeds_pulid(preprocessor_output)
+        elif self.is_instantid:
             return self._get_image_embeds_instantid(preprocessor_output)
         elif self.is_faceid and self.is_plus:
             # Note: FaceID plus uses both face_embed and clip_embed.