From 464889529132d2a6921dfb8fe80abe7db2f02d38 Mon Sep 17 00:00:00 2001
From: Kayla Man <moonyuet2008@hotmail.com>
Date: Fri, 15 Dec 2023 23:09:11 +0800
Subject: [PATCH] make sure the contextlib.nested used before material loading
 while it is compatible for both python2 and 3

---
 openpype/hosts/maya/api/lib.py                |  5 ++-
 .../maya/plugins/publish/extract_playblast.py | 42 ++++++++++++-------
 .../maya/plugins/publish/extract_thumbnail.py | 21 +++++-----
 3 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/openpype/hosts/maya/api/lib.py b/openpype/hosts/maya/api/lib.py
index 817688258b..41290b805e 100644
--- a/openpype/hosts/maya/api/lib.py
+++ b/openpype/hosts/maya/api/lib.py
@@ -172,8 +172,9 @@ def maintained_selection():
             cmds.select(clear=True)
 
 
-def reload_textures():
-    """Reload textures during playblast
+def reload_all_udim_tile_previews():
+    """Regenerate all UDIM tile preview in texture file
+    nodes during context
     """
     texture_files = cmds.ls(type="file")
     if texture_files:
diff --git a/openpype/hosts/maya/plugins/publish/extract_playblast.py b/openpype/hosts/maya/plugins/publish/extract_playblast.py
index b540a2c56d..7bcddf97f1 100644
--- a/openpype/hosts/maya/plugins/publish/extract_playblast.py
+++ b/openpype/hosts/maya/plugins/publish/extract_playblast.py
@@ -43,17 +43,13 @@ class ExtractPlayblast(publish.Extractor):
                     json.dumps(preset, indent=4, sort_keys=True)
                 )
             )
-        if "textures" in preset["viewport_options"] and (
-            "reloadTextures" in preset["viewport_options"]
-        ):
-            with lib.material_loading_mode():
-                lib.reload_textures()
-                # not supported by `capture`
-                preset["viewport_options"].pop("reloadTextures", None)
-                path = capture.capture(log=self.log, **preset)
-        else:
-            preset["viewport_options"].pop("reloadTextures", None)
-            path = capture.capture(log=self.log, **preset)
+
+        if preset["viewport_options"].get("reloadTextures"):
+            # Regenerate all UDIM tiles previews
+            lib.reload_all_udim_tile_previews()
+            # not supported by `capture`
+        preset["viewport_options"].pop("reloadTextures", None)
+        path = capture.capture(log=self.log, **preset)
         self.log.debug("playblast path  {}".format(path))
 
     def process(self, instance):
@@ -206,11 +202,23 @@ class ExtractPlayblast(publish.Extractor):
         # TODO: Remove once dropping Python 2.
         if getattr(contextlib, "nested", None):
             # Python 3 compatibility.
-            with contextlib.nested(
-                lib.maintained_time(),
-                panel_camera(instance.data["panel"], preset["camera"])
-            ):
-                self._capture(preset)
+            if preset["viewport_options"].get("textures"):
+                # If capture includes textures then ensure material
+                # load mode is set to `immediate` to ensure all
+                # textures have loaded when playblast starts
+                with contextlib.nested(
+                    lib.maintained_time(),
+                    panel_camera(instance.data["panel"], preset["camera"]),
+                    lib.material_loading_mode()
+                ):
+                    self._capture(preset)
+
+            else:
+                with contextlib.nested(
+                    lib.maintained_time(),
+                    panel_camera(instance.data["panel"], preset["camera"])
+                ):
+                    self._capture(preset)
         else:
             # Python 2 compatibility.
             with contextlib.ExitStack() as stack:
@@ -218,6 +226,8 @@ class ExtractPlayblast(publish.Extractor):
                 stack.enter_context(
                     panel_camera(instance.data["panel"], preset["camera"])
                 )
+                if preset["viewport_options"].get("textures"):
+                    stack.enter_context(lib.material_loading_mode())
 
                 self._capture(preset)
 
diff --git a/openpype/hosts/maya/plugins/publish/extract_thumbnail.py b/openpype/hosts/maya/plugins/publish/extract_thumbnail.py
index 10082436d6..b24cda8f07 100644
--- a/openpype/hosts/maya/plugins/publish/extract_thumbnail.py
+++ b/openpype/hosts/maya/plugins/publish/extract_thumbnail.py
@@ -152,19 +152,18 @@ class ExtractThumbnail(publish.Extractor):
                         json.dumps(preset, indent=4, sort_keys=True)
                     )
                 )
+
+            if "reloadTextures" in preset["viewport_options"]:
+                lib.reload_all_udim_tile_previews()
+
+            preset["viewport_options"].pop("reloadTextures", None)
             if "textures" in preset["viewport_options"]:
-                if "reloadTextures" in preset["viewport_options"]:
-                    with lib.material_loading_mode():
-                        lib.reload_textures()
-                        preset["viewport_options"].pop("reloadTextures", None)
-                        path = capture.capture(**preset)
-                else:
-                    self.log.debug(
-                        "Reload Textures during playblasting is disabled.")
-                    preset["viewport_options"].pop("reloadTextures", None)
+                with lib.material_loading_mode():
                     path = capture.capture(**preset)
-            # not supported by `capture`
-            path = capture.capture(**preset)
+            else:
+                self.log.debug("Reload Textures during playblasting is disabled.")
+                path = capture.capture(**preset)
+
             playblast = self._fix_playblast_output_path(path)
 
         _, thumbnail = os.path.split(playblast)