refactor seems sound but emitting a cuda error now

Author: dmarx

src/pytti/LossAug/LossOrchestratorClass.py

@@ -127,43 +127,47 @@ def configure_stabilization_augs(img, init_image_pil, params, loss_augs):
 
 
 def configure_optical_flows(img, params, loss_augs):
+    logger.debug(params.device)
+    _device = params.device
+    optical_flows = []
     if params.animation_mode == "Video Source":
         if params.flow_stabilization_weight == "":
             params.flow_stabilization_weight = "0"
-        # if flow stabilization weight is 0, shouldn't this next block just get skipped?
+        # TODO: if flow stabilization weight is 0, shouldn't this next block just get skipped?
 
         for i in range(params.flow_long_term_samples + 1):
-            name = f"optical flow stabilization (frame {-2**i})"
-            weight = params.flow_stabilization_weight
-            comp = torch.zeros(1, 1, 1, 1)  # ,device=device)
             optical_flow = OpticalFlowLoss(
-                comp=comp,
-                weight=weight,
-                name=f"{name} (direct)",
+                comp=torch.zeros(1, 1, 1, 1, device=_device),  # ,device=DEVICE)
+                weight=params.flow_stabilization_weight,
+                name=f"optical flow stabilization (frame {-2**i}) (direct)",
                 image_shape=img.image_shape,
+                device=_device,
             )  # , device=device)
             optical_flow.set_enabled(False)
-            loss_augs.append(optical_flow)
+            optical_flows.append(optical_flow)
 
     elif params.animation_mode == "3D" and params.flow_stabilization_weight not in [
         "0",
         "",
     ]:
-        optical_flows = [
-            TargetFlowLoss.TargetImage(
-                f"optical flow stabilization:{params.flow_stabilization_weight}",
-                img.image_shape,
-                device="cuda",
-            )
-        ]
-        for optical_flow in optical_flows:
-            optical_flow.set_enabled(False)
-        loss_augs.extend(optical_flows)
-    else:
-        optical_flows = []
+        optical_flow = TargetFlowLoss(
+            comp=torch.zeros(1, 1, 1, 1, device=_device),
+            weight=params.flow_stabilization_weight,
+            name="optical flow stabilization (direct)",
+            image_shape=img.image_shape,
+            device=_device,
+        )
+        optical_flow.set_enabled(False)
+        optical_flows.append(optical_flow)
+
+    loss_augs.extend(optical_flows)
+
+    # this shouldn't be in this function based on the name.
     # other loss augs
     if params.smoothing_weight != 0:
-        loss_augs.append(TVLoss(weight=params.smoothing_weight))
+        loss_augs.append(
+            TVLoss(weight=params.smoothing_weight)
+        )  # , device=params.device))
 
     return img, loss_augs, optical_flows
 

tests/test_loss_refactoring.py

@@ -82,3 +82,16 @@ def test_video_optical_flow():
 device: '{TEST_DEVICE}'
 """
     run_cfg(cfg_str)
+
+
+def test_3D_optical_flow():
+    cfg_str = f"""# @package _global_
+scenes: a photograph of an apple
+animation_mode: 3D
+video_path: {video_fpath}
+flow_stabilization_weight: 1
+steps_per_frame: 10
+steps_per_scene: 150
+device: '{TEST_DEVICE}'
+"""
+    run_cfg(cfg_str)