Reorganize bisenet code location, add comments, remove dead code

{models/BiSeNet → BiSeNet}/evaluate.py

@@ -2,7 +2,7 @@
 # -*- encoding: utf-8 -*-
 
 from logger import setup_logger
-from model import BiSeNet
+import BiSeNet
 from face_dataset import FaceMask
 
 import torch

{models/BiSeNet → BiSeNet}/model.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torchvision
 
-from models.BiSeNet.resnet import Resnet18
+from BiSeNet.resnet import Resnet18
 # from modules.bn import InPlaceABNSync as BatchNorm2d
 
 

{models/BiSeNet → BiSeNet}/test.py

@@ -2,7 +2,7 @@
 # -*- encoding: utf-8 -*-
 
 from logger import setup_logger
-from model import BiSeNet
+import BiSeNet
 
 import torch
 

{models/BiSeNet → BiSeNet}/train.py

@@ -1,8 +1,8 @@
 #!/usr/bin/python
 # -*- encoding: utf-8 -*-
 
-from logger import setup_logger
-from model import BiSeNet
+from logger import setup_logger 
+import BiSeNet
 from face_dataset import FaceMask
 from loss import OhemCELoss
 from evaluate import evaluate

app.py

@@ -5,26 +5,22 @@ import glob
 import spaces
 import numpy as np
 
-from datetime import datetime
 from PIL import Image
 from diffusers.utils import load_image
 from diffusers import EulerDiscreteScheduler
 from pipline_ConsistentID import ConsistentIDPipeline
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument('--base_model_path', type=str, 
+                    default="SG161222/Realistic_Vision_V4.0_noVAE")
+parser.add_argument('--gpu', type=int, default=0)
+args = parser.parse_args()
 
-# zero = torch.Tensor([0]).cuda()
-# print(zero.device) # <-- 'cpu' 🤔
-# device = zero.device # "cuda"
-device = "cuda"
-
-# Gets the absolute path of the current script
-script_directory = os.path.dirname(os.path.realpath(__file__))
-
-# download ConsistentID checkpoint to cache
-base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
+device = f"cuda:{args.gpu}"
 
 ### Load base model
 pipe = ConsistentIDPipeline.from_pretrained(
-    base_model_path, 
+    args.base_model_path, 
     torch_dtype=torch.float16, 
     safety_checker=None, # use_safetensors=True, 
     variant="fp16"
@@ -33,11 +29,10 @@ pipe = ConsistentIDPipeline.from_pretrained(
 ### Load consistentID_model checkpoint
 pipe.load_ConsistentID_model(
     consistentID_weight_path="./models/ConsistentID-v1.bin",
-    bise_net_weight_path="./models/face_parsing.pth",
+    bise_net_weight_path="./models/BiSeNet_pretrained_for_ConsistentID.pth",
 )
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
-
 @spaces.GPU
 def process(selected_template_images, custom_image, prompt, 
             negative_prompt, prompt_selected, model_selected_tab, 
@@ -94,7 +89,6 @@ def process(selected_template_images, custom_image, prompt,
     return np.array(images)
 
 # Gets the templates
-script_directory = os.path.dirname(os.path.realpath(__file__))
 preset_template = glob.glob("./images/templates/*.png")
 preset_template = preset_template + glob.glob("./images/templates/*.jpg")
 

pipline_ConsistentID.py

@@ -12,7 +12,7 @@ from huggingface_hub.utils import validate_hf_hub_args
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
-from diffusers.utils import _get_model_file
+from diffusers.utils.import_utils import is_xformers_available
 from functions import insert_markers_to_prompt, masks_for_unique_values, apply_mask_to_raw_image, tokenize_and_mask_noun_phrases_ends, prepare_image_token_idx
 from functions import ProjPlusModel, masks_for_unique_values
 from attention import Consistent_IPAttProcessor, Consistent_AttProcessor, FacialEncoder
@@ -21,7 +21,7 @@ from huggingface_hub import hf_hub_download
 ### Model can be imported from https://github.com/zllrunning/face-parsing.PyTorch?tab=readme-ov-file
 ### We use the ckpt of 79999_iter.pth: https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812
 ### Thanks for the open source of face-parsing model.
-from models.BiSeNet.model import BiSeNet
+from BiSeNet.model import BiSeNet
 import os
 
 PipelineImageInput = Union[
@@ -310,12 +310,10 @@ class ConsistentIDPipeline(StableDiffusionPipeline):
         max_text_length=330      
         if len(self.tokenizer(prompt_face, max_length=self.tokenizer.model_max_length, 
                               padding="max_length", truncation=False, return_tensors="pt").input_ids[0]) != 77:
+            # Put face_caption_align at the beginning of the prompt, so that the original prompt is truncated,
+            # but the face_caption_align is well kept.
             prompt_face = "Detail: " + face_caption_align + " Caption:" + prompt
-        
-        if len(face_caption) > max_text_length:
-            prompt_face = prompt
-            face_caption_align =  ""
-  
+
         # Remove "<|facial|>" from prompt_face.
         # augmented_prompt: 'A person, police officer, half body shot Detail: 
         # The person has one nose , two ears , two eyes , and a mouth , '