Update app.py

app.py

@@ -228,16 +228,25 @@ def run_rmbg(img, sigma=0.0):
     return result.clip(0, 255).astype(np.uint8), alpha
 
 @torch.inference_mode()
-def merge_alpha(img):
+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:
         return img
-    else:      
+    elif C == 4:        
         rgb = img[:, :, :3]
-        alpha = img[:, :, 3] / 255.0  
-        result = 127 + (rgb.astype(np.float32) - 127 + sigma) * alpha[:, :, np.newaxis]
+        alpha = img[:, :, 3] 
+        result = 127 + (rgb.astype(np.float32) - 127 + 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()
@@ -268,6 +277,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)
@@ -289,7 +299,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(
@@ -391,14 +402,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("## wow dub")
     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', 'RGB'])
+                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,
@@ -413,8 +422,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)
@@ -428,15 +437,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)