add output caching

lock requirements
move assets to git lfs
add comments in code

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.glb filter=lfs diff=lfs merge=lfs -text
+*.ply filter=lfs diff=lfs merge=lfs -text

files/horse.ply.output/output.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6e57814fdca91e80b2dad1912093f94d348802be932e3b0560d73c895519a31
+size 1700664

main.py

@@ -27,6 +27,12 @@ def convert_formats(path_input, target_ext):
 
 
 def add_lights(path_input, path_output):
+    """
+    Adds directional lights in the horizontal plane to the glb file.
+    :param path_input: path to input glb
+    :param path_output: path to output glb
+    :return: None
+    """
     glb = pygltflib.GLTF2().load(path_input)
 
     N = 3  # default max num lights in Babylon.js is 4
@@ -34,11 +40,7 @@ def add_lights(path_input, path_output):
 
     lights_extension = {
         "lights": [
-            {
-                "type": "directional",
-                "color": [1.0, 1.0, 1.0],
-                "intensity": 2.0
-            }
+            {"type": "directional", "color": [1.0, 1.0, 1.0], "intensity": 2.0}
             for _ in range(N)
         ]
     }
@@ -50,19 +52,10 @@ def add_lights(path_input, path_output):
     light_nodes = []
     for i in range(N):
         angle = i * angle_step
-        rotation = [
-            0.0,
-            math.sin(angle / 2),
-            0.0,
-            math.cos(angle / 2)
-        ]
+        rotation = [0.0, math.sin(angle / 2), 0.0, math.cos(angle / 2)]
         node = {
             "rotation": rotation,
-            "extensions": {
-                "KHR_lights_punctual": {
-                    "light": i
-                }
-            }
+            "extensions": {"KHR_lights_punctual": {"light": i}},
         }
         light_nodes.append(node)
 
@@ -71,7 +64,7 @@ def add_lights(path_input, path_output):
 
     root_node_index = glb.scenes[glb.scene].nodes[0]
     root_node = glb.nodes[root_node_index]
-    if hasattr(root_node, 'children'):
+    if hasattr(root_node, "children"):
         root_node.children.extend(light_node_indices)
     else:
         root_node.children = light_node_indices
@@ -85,14 +78,20 @@ class Model3D(gr.Model3D):
     """
 
     def postprocess(self, y: str | Path | None) -> dict[str, str] | None:
+        """
+        Converts user input 3D model in any format acceptable by trimesh to glb, required by the web component
+        :param y: path to input 3D model
+        :return: file name mapped to base64 url data
+        """
         if y is not None:
             y = convert_formats(y, "glb")
         out = super().postprocess(y)
         return out
 
 
-def breathe_new_life_into_3d_model(path_input, prompt):
+def breathe_new_life_into_3d_model(path_input, prompt, path_output=None):
     """
+    Paints the input 3D model using the following method:
     @inproceedings{wang2023breathing,
         title={Breathing New Life into 3D Assets with Generative Repainting},
         author={Wang, Tianfu and Kanakis, Menelaos and Schindler, Konrad and Van Gool, Luc and Obukhov, Anton},
@@ -100,7 +99,17 @@ def breathe_new_life_into_3d_model(path_input, prompt):
         year={2023},
         publisher={BMVA Press}
     }
+    :param path_input: path to input 3D model
+    :param prompt: text description of the expected output
+    :param path_output: path to precomputed (cached) output if available
+    :return: path to the painted 3D model in glb format with lights for Gradio web component
     """
+    if path_output is not None:
+        path_output = path_output.name
+        path_output_glb_vis = path_output[:-4] + "_vis.glb"
+        add_lights(path_output, path_output_glb_vis)
+        return path_output_glb_vis
+
     path_output_dir = path_input + ".output"
     os.makedirs(path_output_dir, exist_ok=True)
 
@@ -129,6 +138,9 @@ def breathe_new_life_into_3d_model(path_input, prompt):
 
 
 def run():
+    """
+    Gradio entry point
+    """
     desc = """
         <p align="center">
         <a title="Website" href="https://www.obukhov.ai/repainting_3d_assets" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
@@ -162,7 +174,10 @@ def run():
                 elem_classes="viewport",
                 label="Input Model",
             ),
-            gr.Textbox(label="Text Prompt"),
+            gr.Textbox(
+                label="Text Prompt",
+            ),
+            gr.File(visible=False, label="Cached Output"),
         ],
         outputs=[
             gr.Model3D(
@@ -175,6 +190,9 @@ def run():
             [
                 os.path.join(os.path.dirname(__file__), "files/horse.ply"),
                 "pastel superhero unicorn",
+                os.path.join(
+                    os.path.dirname(__file__), "files/horse.ply.output/output.glb"
+                ),
             ],
         ],
         cache_examples=True,

requirements.txt

@@ -1,3 +1,4 @@
-gradio
-pygltflib
-trimesh
+gradio==3.44.4
+gradio_client==0.5.1
+pygltflib==1.16.0
+trimesh==3.23.5