Add examples

.gitattributes

@@ -34,3 +34,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 example.png filter=lfs diff=lfs merge=lfs -text
+example/aerial_rocks_02_diff_4k_Displacement.jpg filter=lfs diff=lfs merge=lfs -text
+example/aerial_rocks_02_diff_4k_Normal.jpg filter=lfs diff=lfs merge=lfs -text
+example/aerial_rocks_02_diff_4k_Roughness.jpg filter=lfs diff=lfs merge=lfs -text
+example/aerial_rocks_02_diff_4k.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -277,14 +277,13 @@ def load_models():
 def process_image(img, model):
     """Process a single image through the model."""
     img = np.array(img).astype(np.float32) / 255.0
-    if len(img.shape) == 2:  # Grayscale
-        img = np.stack([img, img, img], axis=-1)
-    if img.shape[2] == 4:  # RGBA
-        img = img[:, :, :3]
-    img = torch.from_numpy(np.transpose(img.copy(), (2, 0, 1))).float().unsqueeze(0)
-    output = model(img).squeeze(0).clamp_(0, 1).detach().numpy()
-    output = np.transpose(output, (1, 2, 0))  # CHW to HWC
-    return (output * 255).astype(np.uint8)
+    if len(img.shape) == 2: img = np.stack([img, img, img], axis=-1)
+    if img.shape[2] == 4: img = img[:, :, :3]
+    tensor = torch.from_numpy(np.transpose(img.copy(), (2, 0, 1))).float().unsqueeze(0)
+    with torch.inference_mode():
+        output = model(tensor).squeeze(0).clamp_(0, 1).numpy()
+    del tensor  # Explicit cleanup
+    return (np.transpose(output, (1, 2, 0)) * 255).astype(np.uint8)
 
 def generate_maps(input_image):
     """Generate Normal, Roughness, and Displacement maps from input diffuse texture."""
@@ -318,4 +317,4 @@ if __name__ == "__main__":
                         rough_out = gr.Image(label="Roughness", height=180)
                         disp_out = gr.Image(label="Displacement", height=180)
             btn.click(fn=generate_maps, inputs=input_img, outputs=[normal_out, rough_out, disp_out])
-        demo.launch(theme=gr.themes.Soft())
+        demo.launch(theme=gr.themes.Soft(), ssr_mode=False)