Update app.py

app.py

@@ -227,6 +227,18 @@ def run_rmbg(img, sigma=0.0):
     result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
     return result.clip(0, 255).astype(np.uint8), alpha
 
+@torch.inference_mode()
+def merge_alpha(img):
+    H, W, C = img.shape
+    print(f"img.shape: {img.shape}")
+    if C == 3:
+        return img
+    else:      
+        rgb = img[:, :, :3]
+        alpha = img[:, :, 3] / 255.0  
+        result = 127 + (rgb.astype(np.float32) - 127 + sigma) * alpha[:, :, np.newaxis]
+        return np.clip(result, 0, 255).astype(np.uint8)
+
 
 @torch.inference_mode()
 def process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
@@ -337,6 +349,7 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
 @torch.inference_mode()
 def process_relight(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
     #input_fg, matting = run_rmbg(input_fg)
+    input_fg = merge_alpha(input_fg)
     results = process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source)
     return input_fg, results