Experimental features and bug fixes

toolkit/models/diffusion_feature_extraction.py

@@ -250,13 +250,14 @@ class DiffusionFeatureExtractor3(nn.Module):
         bs = noise_pred.shape[0]
         noise_pred_chunks = torch.chunk(noise_pred, bs)
         timestep_chunks = torch.chunk(timesteps, bs)
+        noisy_latent_chunks = torch.chunk(noisy_latents, bs)
         stepped_chunks = []
         for idx in range(bs):
             model_output = noise_pred_chunks[idx]
             timestep = timestep_chunks[idx]
             scheduler._step_index = None
             scheduler._init_step_index(timestep)
-            sample = noisy_latents.to(torch.float32)
+            sample = noisy_latent_chunks[idx].to(torch.float32)
             
             sigma = scheduler.sigmas[scheduler.step_index]
             sigma_next = scheduler.sigmas[-1] # use last sigma for final step

toolkit/models/flux.py

@@ -117,16 +117,18 @@ def split_gpu_single_block_forward(
     return hidden_state_out
 
 
-def add_model_gpu_splitter_to_flux(transformer: FluxTransformer2DModel):
+def add_model_gpu_splitter_to_flux(
+    transformer: FluxTransformer2DModel,
+    # ~ 5 billion for all other params
+    other_module_params: Optional[int] = 5e9,
+    # since they are not trainable, multiply by smaller number
+    other_module_param_count_scale: Optional[float] = 0.3
+):
     gpu_id_list = [i for i in range(torch.cuda.device_count())]
     
     # if len(gpu_id_list) > 2:
     #     raise ValueError("Cannot split to more than 2 GPUs currently.")
-    
-    # ~ 5 billion for all other params
-    other_module_params = 5e9
-    # since they are not trainable, multiply by smaller number
-    other_module_params *= 0.5
+    other_module_params *= other_module_param_count_scale
     
     # since we are not tuning the 
     total_params = sum(p.numel() for p in transformer.parameters()) + other_module_params