Ensure AE on GPU

app.py

@@ -23,20 +23,17 @@ from pulid.utils import resize_numpy_image_long, seed_everything
 def get_models(name: str, device: torch.device, offload: bool, fp8: bool):
     t5 = load_t5(device, max_length=128)
     clip = load_clip(device)
-    if fp8:
-        model = load_flow_model_quintized(name, device="cpu" if offload else device)
-    else:
-        model = load_flow_model(name, device="cpu" if offload else device)
+    model = load_flow_model(name, device="cpu" if offload else device)
     model.eval()
-    ae = load_ae(name, device="cpu" if offload else device)
-    ae.to(dtype=torch.bfloat16)
+    ae = load_ae(name, device=device)
+    ae.to(device=device, dtype=torch.bfloat16)
     return model, ae, t5, clip
 
 
 class FluxGenerator:
     def __init__(self, model_name: str, device: str, offload: bool, aggressive_offload: bool, args):
         self.device = torch.device(device)
-        self.offload = offload
+        self.offload = False
         self.aggressive_offload = aggressive_offload
         self.model_name = model_name
         self.model, self.ae, self.t5, self.clip_model = get_models(
@@ -45,13 +42,7 @@ class FluxGenerator:
             offload=self.offload,
             fp8=args.fp8,
         )
-        self.pulid_model = PuLIDPipeline(self.model, device="cpu" if offload else device, weight_dtype=torch.bfloat16,
-                                         onnx_provider=args.onnx_provider)
-        if offload:
-            self.pulid_model.face_helper.face_det.mean_tensor = self.pulid_model.face_helper.face_det.mean_tensor.to(torch.device("cuda"))
-            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 = PuLIDPipeline(self.model, device='cuda', weight_dtype=torch.bfloat16)
         self.pulid_model.load_pretrain(args.pretrained_model)
 
     # function to encode an image into latents
@@ -458,5 +449,8 @@ if __name__ == "__main__":
     if args.aggressive_offload:
         args.offload = True
 
+    print(f"Using device: {args.device}")
+    print(f"Offload: {args.offload}")
+
     demo = create_demo(args, args.name, args.device, args.offload, args.aggressive_offload)
     demo.launch(ssr_mode=False)