fix some bug in RVC

trellis/pipelines/samplers/flow_euler.py

@@ -183,6 +183,8 @@ class FlowEulerSampler(Sampler):
         model,
         slat_decoder_gs,
         slat_decoder_mesh,
+        std, 
+        mean,
         dreamsim_model,
         learning_rate,
         input_images,
@@ -222,10 +224,10 @@ class FlowEulerSampler(Sampler):
             for step in range(total_steps):
                 optimizer.zero_grad()
                 pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v_opt)
-                pred_gs = slat_decoder_gs(pred_x_0)
-                # pred_mesh = slat_decoder_mesh(pred_x_0)
+                pred_gs = slat_decoder_gs(pred_x_0 * std + mean)
+                # pred_mesh = slat_decoder_mesh(pred_x_0 * std + mean)
                 rend_gs = render_utils.render_frames(pred_gs[0], extrinsics, intrinsics, {'resolution': 259, 'bg_color': (0, 0, 0)}, need_depth=True, opt=True)['color']
-                # rend_mesh = render_utils.render_frames_opt(pred_mesh[0], extrinsics, intrinsics, {'resolution': 518, 'bg_color': (0, 0, 0)}, need_depth=True)['color']
+                # rend_mesh = render_utils.render_frames_opt(pred_mesh[0], extrinsics, intrinsics, {'resolution': 518, 'bg_color': (0, 0, 0)}, need_depth=True, opt=True)['color']
                 rend_gs = torch.stack(rend_gs, dim=0)
                 loss_gs = loss_utils.l1_loss(rend_gs, input_images, size_average=False).mean(dim=(1,2,3)) + \
                     (1 - loss_utils.ssim(rend_gs, input_images, size_average=False)) + \
@@ -345,6 +347,8 @@ class FlowEulerSampler(Sampler):
         model,
         slat_decoder_gs,
         slat_decoder_mesh,
+        std,
+        mean,
         dreamsim_model,
         apperance_learning_rate,
         start_t,
@@ -392,7 +396,7 @@ class FlowEulerSampler(Sampler):
             else:
                 # learning_rate = cosine_anealing(i - int(np.where(t_seq <= start_t)[0].min()), int(steps - np.where(t_seq <= start_t)[0].min()), apperance_learning_rate, 1e-5)
                 learning_rate = apperance_learning_rate
-                out = self.sample_slat_once_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, dreamsim_model, learning_rate, input_images, extrinsics, intrinsics, sample, t, t_prev, cond, **kwargs)
+                out = self.sample_slat_once_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, std, mean, dreamsim_model, learning_rate, input_images, extrinsics, intrinsics, sample, t, t_prev, cond, **kwargs)
                 sample = out.pred_x_prev
                 ret.pred_x_t.append(out.pred_x_prev)
                 ret.pred_x_0.append(out.pred_x_0)
@@ -865,6 +869,8 @@ class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSa
         model,
         slat_decoder_gs,
         slat_decoder_mesh,
+        std,
+        mean,
         dreamsim_model,
         apperance_learning_rate,
         start_t,
@@ -902,7 +908,7 @@ class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSa
             - 'pred_x_t': a list of prediction of x_t.
             - 'pred_x_0': a list of prediction of x_0.
         """
-        return super().sample_slat_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, dreamsim_model, apperance_learning_rate, start_t, input_images, extrinsics, intrinsics,noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
+        return super().sample_slat_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, std, mean, dreamsim_model, apperance_learning_rate, start_t, input_images, extrinsics, intrinsics,noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
 
 
 class LatentMatchGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, LatentMatchSampler):