import gradio as gr
import torch
import spaces
from diffusers import FluxInpaintPipeline
from PIL import Image, ImageFile

#ImageFile.LOAD_TRUNCATED_IMAGES = True

# Initialize the pipeline
pipe = FluxInpaintPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")
pipe.load_lora_weights(
    "ali-vilab/In-Context-LoRA", 
    weight_name="visual-identity-design.safetensors"
)

def square_center_crop(img, target_size=768):
    if img.mode in ('RGBA', 'P'):
        img = img.convert('RGB')

    width, height = img.size
    crop_size = min(width, height)

    left = (width - crop_size) // 2
    top = (height - crop_size) // 2
    right = left + crop_size
    bottom = top + crop_size

    img_cropped = img.crop((left, top, right, bottom))
    return img_cropped.resize((target_size, target_size), Image.Resampling.LANCZOS)

def duplicate_horizontally(img):
    width, height = img.size
    if width != height:
        raise ValueError(f"Input image must be square, got {width}x{height}")

    new_image = Image.new('RGB', (width * 2, height))
    new_image.paste(img, (0, 0))
    new_image.paste(img, (width, 0))
    return new_image

# Load the mask image
mask = Image.open("mask_square.png")

@spaces.GPU
def generate(image, prompt_user, progress=gr.Progress(track_tqdm=True)):
    prompt_structure = "The two-panel image showcases the logo of a brand, [LEFT] the left panel is showing the logo [RIGHT] the right panel has this logo applied to "
    prompt = prompt_structure + prompt_user

    cropped_image = square_center_crop(image)
    logo_dupli = duplicate_horizontally(cropped_image)

    out = pipe(
        prompt=prompt,
        image=logo_dupli,
        mask_image=mask,
        guidance_scale=6,
        height=768,
        width=1536,
        num_inference_steps=28,
        max_sequence_length=256,
        strength=1
    ).images[0]

    width, height = out.size
    half_width = width // 2
    image_2 = out.crop((half_width, 0, width, height))
    return image_2, out

with gr.Blocks() as demo:
    gr.Markdown("# Logo in Context")
    gr.Markdown("### In-Context LoRA + Image-to-Image, apply your logo to anything")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Logo Image",
                type="pil",
                height=384
            )
            prompt_input = gr.Textbox(
                label="Where should the logo be applied?",
                placeholder="e.g., a coffee cup on a wooden table",
                lines=2
            )
            generate_btn = gr.Button("Generate Application", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Generated Application")
            output_side = gr.Image(label="Side by side")
    with gr.Row():
        gr.Markdown("""
        ### Instructions:
        1. Upload a logo image (preferably square)
        2. Describe where you'd like to see the logo applied
        3. Click 'Generate Application' and wait for the result

        Note: The generation process might take a few moments.
        """)

    # Set up the click event
    generate_btn.click(
        fn=generate,
        inputs=[input_image, prompt_input],
        outputs=[output_image, output_side]
    )

demo.launch()