import gradio as gr
import numpy as np


def custom_composite(im):
    # Get image dimensions and check the number of channels
    h, w, c = im["background"].shape  # c could be 3 (RGB) or 4 (RGBA)

    # Create a black background with 3 channels (RGB)
    black_bg = np.zeros((h, w, 3), dtype=np.uint8)

    # Extract layers, ensuring they have 3 channels (RGB)
    def to_rgb(layer):
        return (
            layer[:, :, :3] if layer.shape[2] == 4 else layer
        )  # Drop alpha if present

    layer1 = to_rgb(im["layers"][0]) if len(im["layers"]) > 0 else black_bg
    layer2 = to_rgb(im["layers"][1]) if len(im["layers"]) > 1 else black_bg

    # Merge layers onto the black background
    final_composite = np.where(layer1 > 0, layer1, black_bg)
    final_composite = np.where(layer2 > 0, layer2, final_composite)

    return final_composite


with gr.Blocks() as demo:
    with gr.Row():
        im = gr.ImageEditor(
            type="numpy", crop_size="1:1", brush=gr.Brush(default_color="white")
        )
        im_preview = gr.Image(format="png")

    im.change(custom_composite, outputs=im_preview, inputs=im, show_progress="hidden")

if __name__ == "__main__":
    demo.launch()