Move to models to GPU inside generate_image

app.py

@@ -20,21 +20,11 @@ from pulid.pipeline_flux import PuLIDPipeline
 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)
-#     model = load_flow_model(name, device="cpu" if offload else device)
-#     model.eval()
-#     ae = load_ae(name, device=device)
-#     return model, ae, t5, clip
-class Tmp:
-    def __init__(self):
-        self.max_length = 128
-
 def get_models(name: str, device: torch.device, offload: bool, fp8: bool):
-    t5 = Tmp()
-    clip = None
-    model = None
+    t5 = load_t5(device, max_length=128)
+    clip = load_clip(device)
+    model = load_flow_model(name, device="cpu" if offload else device)
+    model.eval()
     ae = load_ae(name, device=device)
     return model, ae, t5, clip
 
@@ -50,8 +40,8 @@ class FluxGenerator:
             offload=self.offload,
             fp8=args.fp8,
         )
-        # self.pulid_model = PuLIDPipeline(self.model, device='cuda', weight_dtype=torch.bfloat16)
-        # self.pulid_model.load_pretrain(args.pretrained_model)
+        self.pulid_model = PuLIDPipeline(self.model, device='cuda', weight_dtype=torch.bfloat16)
+        self.pulid_model.load_pretrain(args.pretrained_model)
 
     @spaces.GPU(duration=30)
     @torch.inference_mode()
@@ -82,6 +72,10 @@ class FluxGenerator:
         """
         Core function that performs the image generation.
         """
+        self.t5.to(self.device)
+        self.clip_model.to(self.device)
+        self.ae.to(self.device)
+        self.model.to(self.device)
         self.t5.max_length = max_sequence_length
 
         # If seed == -1, random
@@ -266,6 +260,7 @@ class FluxGenerator:
             )
 
         # Offload flux model, load auto-decoder
+        self.ae.decoder.to(self.device)
         if self.offload:
             self.model.cpu()
             torch.cuda.empty_cache()