fix rmbg for hunyuan

This commit is contained in:
Denis Kochetov
2025-07-17 18:53:42 +00:00
parent aedee80e37
commit 78690f681e
3 changed files with 43 additions and 9 deletions
@@ -14,12 +14,25 @@
from PIL import Image
from rembg import remove, new_session
import numpy as np
class BackgroundRemover():
def __init__(self):
self.session = new_session()
def __init__(self, model_name='u2net'):
self.session = new_session(model_name)
def __call__(self, image: Image.Image):
output = remove(image, session=self.session, bgcolor=[255, 255, 255, 0])
if image.mode in ('RGBA', 'LA') or image.mode != 'RGB':
background = Image.new('RGB', image.size, (255, 255, 255))
if image.mode in ('RGBA', 'LA'):
background.paste(image, mask=image.split()[-1])
else:
background.paste(image)
image = background
output = remove(image, session=self.session)
if output.mode != 'RGBA':
output = output.convert('RGBA')
return output
+16 -3
View File
@@ -14,12 +14,25 @@
from PIL import Image
from rembg import remove, new_session
import numpy as np
class BackgroundRemover():
def __init__(self):
self.session = new_session()
def __init__(self, model_name='u2net'):
self.session = new_session(model_name)
def __call__(self, image: Image.Image):
output = remove(image, session=self.session, bgcolor=[255, 255, 255, 0])
if image.mode in ('RGBA', 'LA') or image.mode != 'RGB':
background = Image.new('RGB', image.size, (255, 255, 255))
if image.mode in ('RGBA', 'LA'):
background.paste(image, mask=image.split()[-1])
else:
background.paste(image)
image = background
output = remove(image, session=self.session)
if output.mode != 'RGBA':
output = output.convert('RGBA')
return output
+11 -3
View File
@@ -4484,12 +4484,20 @@ class Multi_Background_Remover:
images = []
for key, tensor_img in mv_inputs.items():
pil_img = torch_imgs_to_pils(tensor_img)[0].convert("RGBA")
if pil_img.mode == "RGB":
pil_img = torch_imgs_to_pils(tensor_img)[0]
if pil_img.mode != "RGBA":
pil_img = rmbg(pil_img.convert("RGB"))
else:
alpha = pil_img.getchannel('A')
if alpha.getextrema()[0] == 255:
rgb_img = Image.new('RGB', pil_img.size, (255, 255, 255))
rgb_img.paste(pil_img, mask=pil_img.split()[-1] if len(pil_img.split()) == 4 else None)
pil_img = rmbg(rgb_img)
images.append(pil_img)
return (images,)
return (images,)
class Hunyuan3D_V2_ShapeGen_MV:
"""