Adjustments for spaces

app.py

@@ -51,7 +51,7 @@ class FluxGenerator:
             self.pulid_model.face_helper.face_det.device = torch.device("cuda")
             self.pulid_model.face_helper.device = torch.device("cuda")
             self.pulid_model.device = torch.device("cuda")
-        self.pulid_model.load_pretrain(args.pretrained_model, version=args.version)
+        self.pulid_model.load_pretrain(args.pretrained_model)
 
     # function to encode an image into latents
     def encode_image_to_latents(self, img, opts):
@@ -453,9 +453,9 @@ if __name__ == "__main__":
     parser.add_argument("--pretrained_model", type=str, help='for development')
     args = parser.parse_args()
 
-    args.fp8 = True
+    # args.fp8 = True
     if args.aggressive_offload:
         args.offload = True
 
-    demo = create_demo(args, args.name, args.device, True, args.aggressive_offload)
+    demo = create_demo(args, args.name, args.device, args.offload, args.aggressive_offload)
     demo.launch(ssr_mode=False)

pulid/pipeline_flux.py

@@ -4,6 +4,7 @@ import cv2
 import insightface
 import torch
 import torch.nn as nn
+from basicsr.utils import img2tensor, tensor2img
 from facexlib.parsing import init_parsing_model
 from facexlib.utils.face_restoration_helper import FaceRestoreHelper
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -15,11 +16,10 @@ from torchvision.transforms.functional import normalize, resize
 from eva_clip import create_model_and_transforms
 from eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
 from pulid.encoders_transformer import IDFormer, PerceiverAttentionCA
-from pulid.utils import img2tensor, tensor2img
 
 
 class PuLIDPipeline(nn.Module):
-    def __init__(self, dit, device, weight_dtype=torch.bfloat16, onnx_provider='gpu', *args, **kwargs):
+    def __init__(self, dit, device, weight_dtype=torch.bfloat16, *args, **kwargs):
         super().__init__()
         self.device = device
         self.weight_dtype = weight_dtype
@@ -68,12 +68,11 @@ class PuLIDPipeline(nn.Module):
         self.eva_transform_std = eva_transform_std
         # antelopev2
         snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
-        providers = ['CPUExecutionProvider'] if onnx_provider == 'cpu' \
-            else ['CUDAExecutionProvider', 'CPUExecutionProvider']
-        self.app = FaceAnalysis(name='antelopev2', root='.', providers=providers)
+        self.app = FaceAnalysis(
+            name='antelopev2', root='.', providers=['CPUExecutionProvider']
+        )
         self.app.prepare(ctx_id=0, det_size=(640, 640))
-        self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx',
-                                                            providers=providers)
+        self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx', providers=['CPUExecutionProvider'])
         self.handler_ante.prepare(ctx_id=0)
 
         gc.collect()
@@ -91,9 +90,9 @@ class PuLIDPipeline(nn.Module):
         self.clip_vision_model = self.clip_vision_model.to(device)
         self.pulid_encoder = self.pulid_encoder.to(device)
 
-    def load_pretrain(self, pretrain_path=None, version='v0.9.0'):
-        hf_hub_download('guozinan/PuLID', f'pulid_flux_{version}.safetensors', local_dir='models')
-        ckpt_path = f'models/pulid_flux_{version}.safetensors'
+    def load_pretrain(self, pretrain_path=None):
+        hf_hub_download('guozinan/PuLID', 'pulid_flux_v0.9.1.safetensors', local_dir='models')
+        ckpt_path = 'models/pulid_flux_v0.9.1.safetensors'
         if pretrain_path is not None:
             ckpt_path = pretrain_path
         state_dict = load_file(ckpt_path)
@@ -116,7 +115,6 @@ class PuLIDPipeline(nn.Module):
         x = x.repeat(1, 3, 1, 1)
         return x
 
-    @torch.no_grad()
     def get_id_embedding(self, image, cal_uncond=False):
         """
         Args:
@@ -126,6 +124,8 @@ class PuLIDPipeline(nn.Module):
         self.debug_img_list = []
         image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
         # get antelopev2 embedding
+        # for k in self.app.models.keys():
+        #     self.app.models[k].session.set_providers(['CUDAExecutionProvider'])
         face_info = self.app.get(image_bgr)
         if len(face_info) > 0:
             face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[
@@ -151,6 +151,7 @@ class PuLIDPipeline(nn.Module):
         # incase insightface didn't detect face
         if id_ante_embedding is None:
             print('fail to detect face using insightface, extract embedding on align face')
+            # self.handler_ante.session.set_providers(['CUDAExecutionProvider'])
             id_ante_embedding = self.handler_ante.get_feat(align_face)
 
         id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device, self.weight_dtype)

requirements.txt

@@ -1,22 +1,25 @@
-torch==2.4.0
-torchvision==0.19.0
-diffusers==0.30.0
+diffusers==0.25.0
+torch==2.1.2
+torchvision==0.16.2
 transformers==4.43.3
-optimum-quanto==0.2.4
-gradio==5.20.0
 opencv-python
-httpx>=0.23.3
-timm
 einops
 ftfy
+basicsr
 facexlib
 insightface
-onnxruntime
+onnx==1.13.1
 onnxruntime-gpu
+onnxruntime==1.14.1
 accelerate
+huggingface-hub
+timm
 SentencePiece
+fire
 safetensors
+numpy==1.24.1
+spaces
+torchsde
 xformers
 apex
-torchsde
-spaces
+gradio==5.20.0