optimize mvdream pipeline

mvdream/pipeline_mvdream.py

@@ -499,6 +499,13 @@ class MVDreamPipeline(DiffusionPipeline):
         # Prepare extra step kwargs.
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
+        context = torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames)
+        torch.cat([camera] * multiplier)
+
+        if image is not None:
+            ip = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
+            ip_img = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+
         # Denoising loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -511,17 +518,17 @@ class MVDreamPipeline(DiffusionPipeline):
                 unet_inputs = {
                     'x': latent_model_input,
                     'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
-                    'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
+                    'context': context,
                     'num_frames': actual_num_frames,
-                    'camera': torch.cat([camera] * multiplier),
+                    'camera': camera,
                 }
 
                 if image is not None:
-                    unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
-                    unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+                    unet_inputs['ip'] = ip
+                    unet_inputs['ip_img'] = ip_img
                 
                 # predict the noise residual
-                noise_pred = self.unet.forward(**unet_inputs)
+                noise_pred = self.unet(**unet_inputs)
 
                 # perform guidance
                 if do_classifier_free_guidance: