from PIL import Image
import gradio as gr

def resize_to_same_height(img1, img2):
    target_height = min(img1.height, img2.height)

    def resize(img):
        aspect_ratio = img.width / img.height
        new_width = int(target_height * aspect_ratio)
        return img.resize((new_width, target_height), Image.LANCZOS)

    return resize(img1), resize(img2)

def resize_to_same_width(img1, img2):
    target_width = min(img1.width, img2.width)

    def resize(img):
        aspect_ratio = img.height / img.width
        new_height = int(target_width * aspect_ratio)
        return img.resize((target_width, new_height), Image.LANCZOS)

    return resize(img1), resize(img2)

def natural_collage(img1, img2, layout="Horizontal"):
    img1 = img1.convert("RGB")
    img2 = img2.convert("RGB")

    if layout == "Horizontal":
        img1_resized, img2_resized = resize_to_same_height(img1, img2)
        total_width = img1_resized.width + img2_resized.width
        height = img1_resized.height
        collage = Image.new("RGB", (total_width, height))
        collage.paste(img1_resized, (0, 0))
        collage.paste(img2_resized, (img1_resized.width, 0))
    else:
        img1_resized, img2_resized = resize_to_same_width(img1, img2)
        total_height = img1_resized.height + img2_resized.height
        width = img1_resized.width
        collage = Image.new("RGB", (width, total_height))
        collage.paste(img1_resized, (0, 0))
        collage.paste(img2_resized, (0, img1_resized.height))

    return collage

demo = gr.Interface(
    fn=natural_collage,
    inputs=[
        gr.Image(type="pil", label="First Image"),
        gr.Image(type="pil", label="Second Image"),
        gr.Radio(["Horizontal", "Vertical"], label="Layout")
    ],
    outputs=gr.Image(label="Smart Collage"),
    title="Mobile-style Collage App",
    description="Creates seamless collage like Redmi Film Roll, preserving full image, with no padding or cropping."
)

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