cleaned up a deprecated code

Author: dmarx

src/pytti/LossAug/LossOrchestratorClass.py

@@ -4,7 +4,6 @@
 
 from pytti.image_models import PixelImage, RGBImage
 
-# from pytti.LossAug import build_loss
 from pytti.LossAug import TVLoss, HSVLoss, OpticalFlowLoss, TargetFlowLoss
 from pytti.Perceptor.Prompt import parse_prompt
 from pytti.eval_tools import parse_subprompt
@@ -30,7 +29,6 @@ def build_loss(
     pil_target: Image,
     device=None,
 ):
-    # from pytti.LossAug import LOSS_DICT
     if device is None:
         device = img.device
 
@@ -39,9 +37,7 @@ def build_loss(
         loss = type(img).get_preferred_loss()
     else:
         loss = LOSS_DICT[weight_name]
-    # out = Loss.TargetImage(
-    #    f"{weight_name} {name}:{weight}", img.image_shape, pil_target
-    # )
+
     if pil_target is not None:
         resized = pil_target.resize(img.image_shape, Image.LANCZOS)
         comp = loss.make_comp(resized, device=device)
@@ -172,247 +168,10 @@ def configure_optical_flows(img, params, loss_augs):
     return img, loss_augs, optical_flows
 
 
-#######################################
-
-
-# class LossBuilder:
-
-#     LOSS_DICT = {"edge": EdgeLoss, "depth": DepthLoss}
-
-#     def __init__(self, weight_name, weight, name, img, pil_target):
-#         self.weight_name = weight_name
-#         self.weight = weight
-#         self.name = name
-#         self.img = img
-#         self.pil_target = pil_target
-
-#     # uh.... should the places this is beind used maybe just use Loss.__init__?
-#     # TO DO: let's make this a class attribute on something
-
-#     @property
-#     def weight_category(self):
-#         return self.weight_name.split("_")[0]
-
-#     @property
-#     def loss_factory(self):
-#         weight_name = self.weight_category
-#         if weight_name == "direct":
-#             Loss = type(self.img).get_preferred_loss()
-#         else:
-#             Loss = self.LOSS_DICT[weight_name]
-#         return Loss
-
-#     def build_loss(self) -> Loss:
-#         """
-#         Given a weight name, weight, name, image, and target image, returns a loss object
-
-#         :param weight_name: The name of the loss function
-#         :param weight: The weight of the loss
-#         :param name: The name of the loss function
-#         :param img: The image to be optimized
-#         :param pil_target: The target image
-#         :return: The loss function.
-#         """
-#         Loss = self.loss_factory
-#         out = Loss.TargetImage(
-#             f"{self.weight_category} {self.name}:{self.weight}",
-#             self.img.image_shape,
-#             self.pil_target,
-#         )
-#         out.set_enabled(self.pil_target is not None)
-#         return out
-
-
 def _standardize_null(weight):
     weight = str(weight).strip()
     if weight in ("", "None"):
         weight = "0"
     if float(weight) == 0:
         weight = ""
     return weight
-
-
-# class LossConfigurator:
-#     """
-#     Groups together procedures for initializing losses
-#     """
-
-#     def __init__(
-#         self,
-#         init_image_pil: Image.Image,
-#         restore: bool,
-#         img: PixelImage,
-#         embedder,
-#         prompts,
-#         # params,
-#         ########
-#         direct_image_prompts,
-#         semantic_stabilization_weight,
-#         init_image,
-#         semantic_init_weight,
-#         animation_mode,
-#         flow_stabilization_weight,
-#         flow_long_term_samples,
-#         smoothing_weight,
-#         ###########
-#         direct_init_weight,
-#         direct_stabilization_weight,
-#         depth_stabilization_weight,
-#         edge_stabilization_weight,
-#     ):
-#         self.init_image_pil = init_image_pil
-#         self.img = img
-#         self.embedder = embedder
-#         self.prompts = prompts
-
-#         self.init_augs = []
-#         self.loss_augs = []
-#         self.optical_flows = []
-#         self.last_frame_semantic = None
-#         self.semantic_init_prompt = None
-
-#         # self.params = params
-#         self.restore = restore
-
-#         ### params
-#         self.direct_image_prompts = direct_image_prompts
-#         self.semantic_stabilization_weight = _standardize_null(
-#             semantic_stabilization_weight
-#         )
-#         self.init_image = init_image
-#         self.semantic_init_weight = _standardize_null(semantic_init_weight)
-#         self.animation_mode = animation_mode
-#         self.flow_stabilization_weight = _standardize_null(flow_stabilization_weight)
-#         self.flow_long_term_samples = flow_long_term_samples
-#         self.smoothing_weight = _standardize_null(smoothing_weight)
-
-#         ######
-#         self.direct_init_weight = _standardize_null(direct_init_weight)
-#         self.direct_stabilization_weight = _standardize_null(
-#             direct_stabilization_weight
-#         )
-#         self.depth_stabilization_weight = _standardize_null(depth_stabilization_weight)
-#         self.edge_stabilization_weight = _standardize_null(edge_stabilization_weight)
-
-#     def process_direct_image_prompts(self):
-#         # prompt parsing shouldn't go here.
-#         self.loss_augs.extend(
-#             type(self.img)
-#             .get_preferred_loss()
-#             .TargetImage(p.strip(), self.img.image_shape, is_path=True)
-#             for p in self.direct_image_prompts.split("|")
-#             if p.strip()
-#         )
-
-#     def process_semantic_stabilization(self):
-#         last_frame_pil = self.init_image_pil
-#         if not last_frame_pil:
-#             last_frame_pil = self.img.decode_image()
-#         self.last_frame_semantic = parse_prompt(
-#             self.embedder,
-#             f"stabilization:{self.semantic_stabilization_weight}",
-#             last_frame_pil,
-#         )
-#         self.last_frame_semantic.set_enabled(self.init_image_pil is not None)
-#         for scene in self.prompts:
-#             scene.append(self.last_frame_semantic)
-
-#     def configure_losses(self):
-#         if self.init_image_pil is not None:
-#             self.configure_init_image()
-#         self.process_direct_image_prompts()
-#         if self.semantic_stabilization_weight:
-#             self.process_semantic_stabilization()
-#         self.configure_stabilization_augs()
-#         self.configure_optical_flows()
-#         self.configure_aesthetic_losses()
-
-#         return (
-#             self.loss_augs,
-#             self.init_augs,
-#             self.stabilization_augs,
-#             self.optical_flows,
-#             self.semantic_init_prompt,
-#             self.last_frame_semantic,
-#             self.img,
-#         )
-
-#     def configure_init_image(self):
-
-#         if not self.restore:
-#             # move these logging statements into .encode_image()
-#             logger.info("Encoding image...")
-#             self.img.encode_image(self.init_image_pil)
-#             logger.info("Encoded Image:")
-#             # pretty sure this assumes we're in a notebook
-#             display.display(self.img.decode_image())
-
-#         ## wrap this for the flexibility that the loop is pretending to provide...
-#         # set up init image prompt
-#         if self.direct_init_weight:
-#             init_aug = LossBuilder(
-#                 "direct_init_weight",
-#                 self.direct_init_weight,
-#                 f"init image ({self.init_image})",
-#                 self.img,
-#                 self.init_image_pil,
-#             ).build_loss()
-#             self.loss_augs.append(init_aug)
-#             self.init_augs.append(init_aug)
-
-#         ########
-#         if self.semantic_init_weight:
-#             self.semantic_init_prompt = parse_prompt(
-#                 self.embedder,
-#                 f"init image [{self.init_image}]:{self.semantic_init_weight}",
-#                 self.init_image_pil,
-#             )
-#             self.prompts[0].append(self.semantic_init_prompt)
-
-#     # stabilization
-#     def configure_stabilization_augs(self):
-#         d_augs = {
-#             "direct_stabilization_weight": self.direct_stabilization_weight,
-#             "depth_stabilization_weight": self.depth_stabilization_weight,
-#             "edge_stabilization_weight": self.edge_stabilization_weight,
-#         }
-#         stabilization_augs = [
-#             LossBuilder(
-#                 k, v, "stabilization", self.img, self.init_image_pil
-#             ).build_loss()
-#             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):
-#         optical_flows = None
-
-#         if self.animation_mode == "Video Source":
-#             if self.flow_stabilization_weight == "":
-#                 self.flow_stabilization_weight = "0"
-#             optical_flows = [
-#                 OpticalFlowLoss.TargetImage(
-#                     f"optical flow stabilization (frame {-2**i}):{self.flow_stabilization_weight}",
-#                     self.img.image_shape,
-#                 )
-#                 for i in range(self.flow_long_term_samples + 1)
-#             ]
-
-#         elif self.animation_mode == "3D" and self.flow_stabilization_weight:
-#             optical_flows = [
-#                 TargetFlowLoss.TargetImage(
-#                     f"optical flow stabilization:{self.flow_stabilization_weight}",
-#                     self.img.image_shape,
-#                 )
-#             ]
-
-#         if optical_flows is not None:
-#             for optical_flow in optical_flows:
-#                 optical_flow.set_enabled(False)
-#             self.loss_augs.extend(optical_flows)
-
-#     def configure_aesthetic_losses(self):
-#         if self.smoothing_weight != 0:
-#             self.loss_augs.append(TVLoss(weight=self.smoothing_weight))

src/pytti/LossAug/MSELossClass.py

@@ -4,10 +4,9 @@
 from torch.nn import functional as F
 from pytti.LossAug.BaseLossClass import Loss
 
-# from pytti.Notebook import Rotoscoper
 from pytti.rotoscoper import Rotoscoper
 from pytti import fetch, vram_usage_mode
-from pytti.eval_tools import parse, parse_subprompt
+from pytti.eval_tools import parse_subprompt
 import torch
 
 

src/pytti/LossAug/OpticalFlowLossClass.py

@@ -16,10 +16,8 @@
 import gma
 from gma.core.network import RAFTGMA
 
-# from gma.core.utils import flow_viz
 from gma.core.utils.utils import InputPadder
 
-# from pytti import fetch, to_pil, DEVICE, vram_usage_mode
 from pytti import fetch, vram_usage_mode
 from pytti.LossAug.MSELossClass import MSELoss
 from pytti.rotoscoper import Rotoscoper
@@ -100,7 +98,6 @@ def init_GMA(checkpoint_path=None, device=None):
             args = parser.parse_args([])
 
             # create new OrderedDict that does not contain `module.` prefix
-            # state_dict = torch.load(checkpoint_path)
             state_dict = torch.load(checkpoint_path, map_location=device)
             from collections import OrderedDict
 
@@ -110,18 +107,9 @@ def init_GMA(checkpoint_path=None, device=None):
                     k = k[7:]  # remove `module.`
                 new_state_dict[k] = v
 
-            # GMA = torch.nn.DataParallel(RAFTGMA(args), device_ids=[device])
             GMA = RAFTGMA(args)
-            # GMA = torch.nn.parallel.DistributedDataParallel(RAFTGMA(args).to(device), device_ids=[device])
-            # GMA = RAFTGMA(args)
-            # GMA.load_state_dict(torch.load(checkpoint_path, map_location=device))
-            # GMA.load_state_dict(torch.load(checkpoint_path))
             GMA.load_state_dict(new_state_dict)
             logger.debug("gma state_dict loaded")
-            ###########################
-            # 1. Fix state dict (remove module prefixes)
-            # 2. load state dict into model without DataParallel
-            ###########################
             GMA.to(device)  # redundant?
             GMA.eval()
 
@@ -209,7 +197,6 @@ def get_loss(self, input, img, device=None):
         if device is None:
             device = getattr(self, "device", self.device)
         init_GMA(
-            # "GMA/checkpoints/gma-sintel.pth"
             device=device,
         )  # update this to use model dir from config
         image1 = self.last_step
@@ -220,8 +207,6 @@ def get_loss(self, input, img, device=None):
         logger.debug(device)
         logger.debug((flow.shape, flow.device))
         logger.debug((self.comp.shape, self.comp.device))
-        # logger.debug(GMA.device) # ugh... I bet this is another dataparallel thing.
-        # logger.debug(GMA.module.device)
         flow = flow.to(device, memory_format=torch.channels_last)
         return super().get_loss(TF.resize(flow, self.comp.shape[-2:]), img) / self.mag
 
@@ -232,9 +217,9 @@ class OpticalFlowLoss(MSELoss):
     def motion_edge_map(
         flow_forward,
         flow_backward,
-        img,  # is this even being used anywhere here?
-        border_mode="smear",
-        sampling_mode="bilinear",
+        img,  # unused
+        border_mode="smear",  # unused
+        sampling_mode="bilinear",  # unused
         device=None,
     ):
         """
@@ -325,7 +310,6 @@ def get_flow(image1, image2, device=None):
         """
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
-        # init_GMA("GMA/checkpoints/gma-sintel.pth")
         init_GMA(
             device=device,
         )
@@ -386,11 +370,9 @@ def set_flow(
             )
         logger.debug(device)
         if path is not None:
-            # img = img.clone()
             img = img.clone().to(device)
             if not isinstance(device, torch.device):
                 device = torch.device(device)
-            # logger.debug(device)
             state_dict = torch.load(path, map_location=device)
             img.load_state_dict(state_dict)
 
@@ -412,13 +394,9 @@ def set_flow(
         image1.add_(noise)
         image2.add_(noise)
 
-        # flow_forward = OpticalFlowLoss.get_flow(image1, image2)
-        # flow_backward = OpticalFlowLoss.get_flow(image2, image1)
-        # flow_forward = self.get_flow(image1, image2, device=device)
-        # flow_backward = self.get_flow(image2, image1, device=device)
         flow_forward = OpticalFlowLoss.get_flow(image1, image2, device=device)
         flow_backward = OpticalFlowLoss.get_flow(image2, image1, device=device)
-        unwarped_target_direct = img.decode_tensor()
+        unwarped_target_direct = img.decode_tensor()  # unused
         flow_target_direct = apply_flow(
             img, -flow_backward, border_mode=border_mode, sampling_mode=sampling_mode
         )