improve memory usage for zero GPUs

pipelines/pipeline_flux_infusenet.py

@@ -359,6 +359,11 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
             lora_scale=lora_scale,
         )
 
+        # CPU offload T5, move back controlnet to GPU
+        self.text_encoder_2.cpu()
+        torch.cuda.empty_cache()
+        self.controlnet.to(device)
+
         # 3. Prepare control image
         num_channels_latents = self.transformer.config.in_channels // 4
         if isinstance(self.controlnet, FluxControlNetModel):
@@ -492,11 +497,6 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
             ]
             controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps)
 
-        # CPU offload T5, move back controlnet to GPU
-        self.text_encoder_2.cpu()
-        torch.cuda.empty_cache()
-        self.controlnet.to(device)
-
         # 7. Denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):

pipelines/pipeline_infu_flux.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import math
 import os
 import random
@@ -199,9 +200,9 @@ class InfUFluxPipeline:
         ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu")
         image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
         del ipm_state_dict
-        image_proj_model.to('cuda', torch.bfloat16)
-        image_proj_model.eval()
-        self.image_proj_model_aes = image_proj_model
+        self.image_proj_model_aes = copy.deepcopy(image_proj_model)
+        self.image_proj_model_aes.to('cuda', torch.bfloat16)
+        self.image_proj_model_aes.eval()
         
         image_proj_model = Resampler(
             dim=1280,
@@ -217,9 +218,9 @@ class InfUFluxPipeline:
         ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu")
         image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
         del ipm_state_dict
-        image_proj_model.to('cpu', torch.bfloat16)
-        image_proj_model.eval()
         self.image_proj_model_sim = image_proj_model
+        self.image_proj_model_sim.to('cpu', torch.bfloat16)
+        self.image_proj_model_sim.eval()
 
         self.image_proj_model = self.image_proj_model_aes