From 520c29eeb2b7b2b50ec4c4fb31008878e08fad39 Mon Sep 17 00:00:00 2001
From: David Marx <david.marx84@gmail.com>
Date: Mon, 11 Apr 2022 22:20:36 -0700
Subject: [PATCH] added missing stab augs, patched yaml turning "off" to False

---
 src/pytti/LossAug/LossOrchestratorClass.py | 2 ++
 src/pytti/workhorse.py                     | 9 ++++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/pytti/LossAug/LossOrchestratorClass.py b/src/pytti/LossAug/LossOrchestratorClass.py
index 65aa04b..6e64c4e 100644
--- a/src/pytti/LossAug/LossOrchestratorClass.py
+++ b/src/pytti/LossAug/LossOrchestratorClass.py
@@ -168,6 +168,7 @@ def configure_losses(self):
         return (
             self.loss_augs,
             self.init_augs,
+            self.stabilization_augs,
             self.optical_flows,
             self.semantic_init_prompt,
             self.last_frame_semantic,
@@ -220,6 +221,7 @@ def configure_stabilization_augs(self):
             for k, v in d_augs.items()
             if v
         ]
+        self.stabilization_augs = stabilization_augs
         self.loss_augs.extend(stabilization_augs)
 
     def configure_optical_flows(self):
diff --git a/src/pytti/workhorse.py b/src/pytti/workhorse.py
index 506d079..cdfc981 100644
--- a/src/pytti/workhorse.py
+++ b/src/pytti/workhorse.py
@@ -176,6 +176,11 @@ def load_video_source(
 def _main(cfg: DictConfig):
     # params = OmegaConf.to_container(cfg, resolve=True)
     params = cfg
+
+    # literal "off" in yaml interpreted as False
+    if params.animation_mode == False:
+        params.animation_mode = "off"
+
     logger.debug(params)
     logger.debug(OmegaConf.to_container(cfg, resolve=True))
     latest = -1
@@ -188,7 +193,7 @@ def _main(cfg: DictConfig):
 
     ### Move these into default.yaml
     # @markdown check `restore` to restore from a previous run
-    restore = params.get('restore') or False # @param{type:"boolean"}
+    restore = params.get("restore") or False  # @param{type:"boolean"}
     # @markdown check `reencode` if you are restoring with a modified image or modified image settings
     reencode = False  # @param{type:"boolean"}
     # @markdown which run to restore
@@ -353,6 +358,7 @@ def do_run():
         (
             loss_augs,
             init_augs,
+            stabilization_augs,
             optical_flows,
             semantic_init_prompt,
             last_frame_semantic,
@@ -443,6 +449,7 @@ def do_run():
             video_frames=video_frames,
             # these can be passed in together as the loss orchestrator
             optical_flows=optical_flows,
+            stabilization_augs=stabilization_augs,
             last_frame_semantic=last_frame_semantic,  # fml...
             semantic_init_prompt=semantic_init_prompt,
             init_augs=init_augs,