stage

optimum/exporters/neuron/convert.py

@@ -345,6 +345,11 @@ def export_models(
     failed_models = []
     total_compilation_time = 0
     compile_configs = {}
+    models_and_neuron_configs.pop("text_encoder")
+    models_and_neuron_configs.pop("text_encoder_2")
+    # models_and_neuron_configs.pop('unet')
+    models_and_neuron_configs.pop('vae_encoder')
+    models_and_neuron_configs.pop('vae_decoder')
     for i, model_name in enumerate(models_and_neuron_configs.keys()):
         logger.info(f"***** Compiling {model_name} *****")
         submodel, sub_neuron_config = models_and_neuron_configs[model_name]

optimum/neuron/modeling_diffusion.py

@@ -1058,15 +1058,15 @@ def forward(
         inputs = (sample, timestep, encoder_hidden_states)
         if timestep_cond is not None:
             inputs = inputs + (timestep_cond,)
-        if added_cond_kwargs is not None:
-            text_embeds = added_cond_kwargs.pop("text_embeds", None)
-            time_ids = added_cond_kwargs.pop("time_ids", None)
-            inputs = inputs + (text_embeds, time_ids)
         if mid_block_additional_residual is not None:
             inputs = inputs + (mid_block_additional_residual,)
         if down_block_additional_residuals is not None:
             for idx in range(len(down_block_additional_residuals)):
                 inputs = inputs + (down_block_additional_residuals[idx],)
+        if added_cond_kwargs:
+            text_embeds = added_cond_kwargs.pop("text_embeds", None)
+            time_ids = added_cond_kwargs.pop("time_ids", None)
+            inputs = inputs + (text_embeds, time_ids)
 
         outputs = self.model(*inputs)
         return outputs
@@ -1139,9 +1139,15 @@ def forward(
         controlnet_cond: torch.Tensor,
         conditioning_scale: float = 1.0,
         guess_mode: bool = False,
+        added_cond_kwargs: Optional[Dict] = None,
         return_dict: bool = True,
     ) -> Union["ControlNetOutput", Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]:
+        timestep = timestep.expand((sample.shape[0],)).to(torch.long)
         inputs = (sample, timestep, encoder_hidden_states, controlnet_cond, conditioning_scale)
+        if added_cond_kwargs:
+            text_embeds = added_cond_kwargs.pop("text_embeds", None)
+            time_ids = added_cond_kwargs.pop("time_ids", None)
+            inputs += (text_embeds, time_ids)
         outputs = self.model(*inputs)
 
         if guess_mode: