Initial commit. Converted for ComfyUI

2025-03-09 15:01:22 -07:00
commit a0ff6b6bc9
32 changed files with 4461 additions and 0 deletions
@@ -0,0 +1,77 @@
+"""
+Example workflow for using MatAnyone in ComfyUI
+
+This is a Python representation of a ComfyUI workflow.
+You can use this as a reference for how to connect nodes.
+"""
+
+import torch
+import numpy as np
+from PIL import Image
+import os
+
+# This is a representation of how you would use the node in ComfyUI
+# In the actual ComfyUI interface, you would connect nodes visually
+
+def example_workflow():
+    # Step 1: Load video frames
+    # In ComfyUI, you would use a video loader node
+    # Here we're just creating dummy frames for illustration
+    video_frames = torch.rand(10, 3, 512, 512)  # 10 frames, RGB, 512x512
+    
+    # Step 2: Create a mask for the first frame
+    # In ComfyUI, you would use a mask drawing tool or load a mask
+    mask = torch.zeros(512, 512)
+    # Draw a circle in the center as an example foreground
+    center_x, center_y = 256, 256
+    radius = 100
+    for y in range(512):
+        for x in range(512):
+            if (x - center_x)**2 + (y - center_y)**2 < radius**2:
+                mask[y, x] = 1.0
+    
+    # Step 3: Process with MatAnyone
+    # In ComfyUI, this would be the MatAnyone node
+    from __init__ import MatAnyoneNode
+    matanyone_node = MatAnyoneNode()
+    
+    # Make sure the model is downloaded
+    model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model")
+    os.makedirs(model_dir, exist_ok=True)
+    model_path = os.path.join(model_dir, "matanyone.pth")
+    
+    if not os.path.exists(model_path):
+        print(f"Please download the model from https://github.com/pq-yang/MatAnyone/releases/download/v1.0.0/matanyone.pth")
+        print(f"and place it in {model_dir}")
+        return
+    
+    # Process the video
+    foreground_frames, alpha_frames = matanyone_node.process_video(
+        video_frames=video_frames,
+        mask=mask,
+        warmup_frames=10,
+        erode_kernel=10,
+        dilate_kernel=10,
+        max_size=-1
+    )
+    
+    # Step 4: Use the results
+    # In ComfyUI, you would connect these to other nodes
+    print(f"Foreground frames shape: {foreground_frames.shape}")
+    print(f"Alpha frames shape: {alpha_frames.shape}")
+    
+    # Example: Save the first frame as an image
+    fg_frame = foreground_frames[0].permute(1, 2, 0).cpu().numpy()
+    fg_frame = (fg_frame * 255).astype(np.uint8)
+    fg_image = Image.fromarray(fg_frame)
+    fg_image.save("example_foreground.png")
+    
+    alpha_frame = alpha_frames[0][0].cpu().numpy()
+    alpha_frame = (alpha_frame * 255).astype(np.uint8)
+    alpha_image = Image.fromarray(alpha_frame)
+    alpha_image.save("example_alpha.png")
+    
+    print("Example images saved!")
+
+if __name__ == "__main__":
+    example_workflow()