update

pipelines/pipeline_flux_infusenet.py

@@ -301,7 +301,8 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
         else:
             batch_size = prompt_embeds.shape[0]
 
-        device = self._execution_device
+        # device = self._execution_device
+        device = 'cuda'
         dtype = self.transformer.dtype
 
         # CPU offload controlnet, move back T5 to GPU

pipelines/pipeline_infu_flux.py

@@ -180,6 +180,9 @@ class InfUFluxPipeline:
         pipe.controlnet.cpu()
         pipe.text_encoder_2.cpu()
         torch.cuda.empty_cache()
+        pipe.text_encoder.to('cuda')
+        pipe.transformer.to('cuda')
+        pipe.vae.to('cuda')
         # pipe.enable_model_cpu_offload()
         self.pipe = pipe
 
@@ -219,7 +222,7 @@ class InfUFluxPipeline:
         image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
         del ipm_state_dict
         self.image_proj_model_sim = image_proj_model
-        self.image_proj_model_sim.to('cpu', torch.bfloat16)
+        self.image_proj_model_sim.to('cpu', dtype=torch.bfloat16)
         self.image_proj_model_sim.eval()
 
         self.image_proj_model = self.image_proj_model_aes