Upload rbga_to_white_bg_2.py
Browse files- rbga_to_white_bg_2.py +54 -0
rbga_to_white_bg_2.py
ADDED
|
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
|
| 3 |
+
class RGBA_White_Background_2:
|
| 4 |
+
"""
|
| 5 |
+
Composites an RGBA image over a white background and outputs RGB.
|
| 6 |
+
Assumes input image in ComfyUI tensor format: float32, 0..1, shape [B,H,W,C].
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
@classmethod
|
| 10 |
+
def INPUT_TYPES(cls):
|
| 11 |
+
return {
|
| 12 |
+
"required": {
|
| 13 |
+
"image": ("IMAGE",), # Can be RGBA or RGB; if RGB, it passes through
|
| 14 |
+
}
|
| 15 |
+
}
|
| 16 |
+
|
| 17 |
+
RETURN_TYPES = ("IMAGE",)
|
| 18 |
+
RETURN_NAMES = ("image",)
|
| 19 |
+
FUNCTION = "composite"
|
| 20 |
+
CATEGORY = "image"
|
| 21 |
+
|
| 22 |
+
def composite(self, image):
|
| 23 |
+
# image: [B,H,W,C], float32 0..1
|
| 24 |
+
if not isinstance(image, torch.Tensor):
|
| 25 |
+
raise TypeError("Expected ComfyUI IMAGE as a torch.Tensor")
|
| 26 |
+
|
| 27 |
+
if image.ndim != 4:
|
| 28 |
+
raise ValueError(f"Expected image shape [B,H,W,C], got {tuple(image.shape)}")
|
| 29 |
+
|
| 30 |
+
b, h, w, c = image.shape
|
| 31 |
+
|
| 32 |
+
# If already RGB, just ensure it's sane and return.
|
| 33 |
+
if c == 3:
|
| 34 |
+
return (image.clamp(0.0, 1.0),)
|
| 35 |
+
|
| 36 |
+
if c != 4:
|
| 37 |
+
raise ValueError(f"Expected 3 (RGB) or 4 (RGBA) channels, got {c}")
|
| 38 |
+
|
| 39 |
+
rgb = image[..., :3]
|
| 40 |
+
a = image[..., 3:4] # keep as [B,H,W,1]
|
| 41 |
+
|
| 42 |
+
# Composite: out = rgb * a + white * (1-a) ; white=1
|
| 43 |
+
out = rgb * a + (1.0 - a)
|
| 44 |
+
|
| 45 |
+
return (out.clamp(0.0, 1.0),)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
NODE_CLASS_MAPPINGS = {
|
| 49 |
+
"RGBA_White_Background_2": RGBA_White_Background_2
|
| 50 |
+
}
|
| 51 |
+
|
| 52 |
+
NODE_DISPLAY_NAME_MAPPINGS = {
|
| 53 |
+
"RGBA_White_Background_2": "RGBA to White Background RGB (White Background)"
|
| 54 |
+
}
|