Update app.py

app.py

@@ -227,6 +227,33 @@ 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, sigma=0.0):
+    if img is None:
+        return None    
+
+    if len(img.shape) == 2:        
+        img = np.stack((img,)*3, axis=-1)
+    
+    H, W, C = img.shape
+    print(f"img.shape: {img.shape}")
+    
+    if C == 3:
+        img, _ = run_rmbg(img)
+        return img
+    elif C == 4:        
+        rgb = img[:, :, :3].astype(np.float32)
+        alpha = img[:, :, 3].astype(np.float32) / 255.0  
+
+        result = rgb * alpha[:, :, np.newaxis] + 255 * (1 - alpha[:, :, np.newaxis])
+
+        if sigma != 0:
+            result += sigma * alpha[:, :, np.newaxis]
+
+        return np.clip(result, 0, 255).astype(np.uint8)
+    else:
+        raise ValueError(f"Unexpected number of channels: {C}")
+
 
 @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):
@@ -256,6 +283,7 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
 
     rng = torch.Generator(device=device).manual_seed(int(seed))
 
+    #fg = input_fg
     fg = resize_and_center_crop(input_fg, image_width, image_height)
 
     concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)
@@ -277,7 +305,8 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
             cross_attention_kwargs={'concat_conds': concat_conds},
         ).images.to(vae.dtype) / vae.config.scaling_factor
     else:
-        bg = resize_and_center_crop(input_bg, image_width, image_height)
+        #bg = input_bg
+        bg = resize_and_center_crop(input_bg, image_width, image_height)        
         bg_latent = numpy2pytorch([bg]).to(device=vae.device, dtype=vae.dtype)
         bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
         latents = i2i_pipe(
@@ -333,10 +362,11 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
     return pytorch2numpy(pixels)
 
 
-@spaces.GPU
+@spaces.GPU(duration=240)
 @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, 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
 
@@ -378,14 +408,12 @@ class BGSource(Enum):
 block = gr.Blocks().queue()
 with block:
     with gr.Row():
-        gr.Markdown("## IC-Light (Relighting with Foreground Condition)")
-    with gr.Row():
-        gr.Markdown("See also https://github.com/lllyasviel/IC-Light for background-conditioned model and normal estimation")
+        gr.Markdown("##ICLight without mask")
     with gr.Row():
         with gr.Column():
             with gr.Row():
-                input_fg = gr.Image(sources='upload', type="numpy", label="Image", height=480)
-                output_bg = gr.Image(type="numpy", label="Preprocessed Foreground", height=480)
+                input_fg = gr.Image(sources='upload', type="numpy", label="Image", image_mode='RGBA')
+                output_bg = gr.Image(type="numpy", label="Preprocessed Foreground")
             prompt = gr.Textbox(label="Prompt")
             bg_source = gr.Radio(choices=[e.value for e in BGSource],
                                  value=BGSource.NONE.value,
@@ -400,8 +428,8 @@ with block:
                     seed = gr.Number(label="Seed", value=12345, precision=0)
 
                 with gr.Row():
-                    image_width = gr.Slider(label="Image Width", minimum=256, maximum=1024, value=512, step=64)
-                    image_height = gr.Slider(label="Image Height", minimum=256, maximum=1024, value=640, step=64)
+                    image_width = gr.Slider(label="Image Width", minimum=256, maximum=2048, value=512, step=64)
+                    image_height = gr.Slider(label="Image Height", minimum=256, maximum=2048, value=640, step=64)
 
             with gr.Accordion("Advanced options", open=False):
                 steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)
@@ -415,15 +443,7 @@ with block:
             result_gallery = gr.Gallery(height=832, object_fit='contain', label='Outputs')
     with gr.Row():
         dummy_image_for_outputs = gr.Image(visible=False, label='Result')
-        gr.Examples(
-            fn=lambda *args: [[args[-1]], "imgs/dummy.png"],
-            examples=db_examples.foreground_conditioned_examples,
-            inputs=[
-                input_fg, prompt, bg_source, image_width, image_height, seed, dummy_image_for_outputs
-            ],
-            outputs=[result_gallery, output_bg],
-            run_on_click=True, examples_per_page=1024
-        )
+
     ips = [input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source]
     relight_button.click(fn=process_relight, inputs=ips, outputs=[output_bg, result_gallery])
     example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)