Reorganize code files

app.py

@@ -8,11 +8,11 @@ import numpy as np
 from PIL import Image
 from diffusers.utils import load_image
 from diffusers import EulerDiscreteScheduler
-from pipline_ConsistentID import ConsistentIDPipeline
+from lib.pipline_ConsistentID import ConsistentIDPipeline
 import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument('--base_model_path', type=str, 
-                    default="SG161222/Realistic_Vision_V4.0_noVAE")
+                    default="models/SG161222/Realistic_Vision_V4.0_noVAE")
 parser.add_argument('--gpu', type=int, default=0)
 args = parser.parse_args()
 

{BiSeNet → lib/BiSeNet}/loss.py

@@ -6,9 +6,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-import numpy as np
-
-
 class OhemCELoss(nn.Module):
     def __init__(self, thresh, n_min, ignore_lb=255, *args, **kwargs):
         super(OhemCELoss, self).__init__()

{BiSeNet → lib/BiSeNet}/makeup.py

@@ -1,5 +1,4 @@
 import cv2
-import os
 import numpy as np
 from skimage.filters import gaussian
 

{BiSeNet → lib/BiSeNet}/model.py

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

attention.py → lib/attention.py

@@ -2,8 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.lora import LoRALinearLayer
-from functions import AttentionMLP
-
+from .functions import AttentionMLP
 
 class FuseModule(nn.Module):
     def __init__(self, embed_dim):

pipline_ConsistentID.py → lib/pipline_ConsistentID.py

@@ -13,15 +13,15 @@ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
 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
+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
 from easydict import EasyDict as edict
 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 BiSeNet.model import BiSeNet
+from .BiSeNet.model import BiSeNet
 import os
 
 PipelineImageInput = Union[
@@ -33,7 +33,7 @@ PipelineImageInput = Union[
 
 ### Download the pretrained model from huggingface and put it locally, then place the model in a local directory and specify the directory location.
 class ConsistentIDPipeline(StableDiffusionPipeline):
-    
+
     def cuda(self, dtype=torch.float16, use_xformers=False):
         self.to('cuda', dtype)
         
@@ -364,14 +364,17 @@ class ConsistentIDPipeline(StableDiffusionPipeline):
 
         return parsed_image_parts, facial_masks, key_masked_raw_images_dict
 
-    # Release the unet and vae models to save memory.
-    def release_unet_vae(self):
-        unet = edict()
-        # Only keep the config and in_channels attributes that are used in the pipeline.
-        unet.config = self.unet.config
-        unet.in_channels = self.unet.in_channels
-        self.unet = unet
-        self.vae = None
+    # Release the unet or vae to save memory.
+    def release_components(self, release_unet=False, release_vae=True):
+        if release_unet:
+            unet = edict()
+            # Only keep the config and in_channels attributes that are used in the pipeline.
+            unet.config = self.unet.config
+            unet.in_channels = self.unet.in_channels
+            self.unet = unet
+
+        if release_vae:
+            self.vae = None
 
     # input_subj_image_obj: an Image object.
     def generate_id_prompt_embeds(self, prompt, negative_prompt, input_subj_image_obj, device, calc_uncond=True):