import gradio as gr
import torch
from PIL import Image
import numpy as np
import cv2
from diffusers import StableDiffusionPipeline

# Setup the model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "s3nh/artwork-arcane-stable-diffusion"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe = pipe.to(device)

# Generate T-shirt design function
def generate_tshirt_design(text):
    prompt = f"{text}"
    image = pipe(prompt).images[0]
    return image

# Remove background from the generated design
def remove_background(design_image):
    design_np = np.array(design_image)
    gray = cv2.cvtColor(design_np, cv2.COLOR_BGR2GRAY)
    _, alpha = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
    b, g, r = cv2.split(design_np)
    rgba = [b, g, r, alpha]
    design_np = cv2.merge(rgba, 4)
    return design_np

# Blend design with T-shirt mockup
def blend_design_with_mockup(mockup, design):
    mockup_np = np.array(mockup)
    design_np = np.array(design)

    # Ensure the design is RGBA
    if design_np.shape[2] == 3:
        alpha = np.ones((design_np.shape[0], design_np.shape[1], 1), dtype=design_np.dtype) * 255
        design_np = np.concatenate((design_np, alpha), axis=2)

    design_resized = cv2.resize(design_np, (mockup_np.shape[1] // 4, mockup_np.shape[0] // 4))  # Adjust size as needed

    y_offset = (mockup_np.shape[0] - design_resized.shape[0]) // 2
    x_offset = (mockup_np.shape[1] - design_resized.shape[1]) // 2

    y1, y2 = y_offset, y_offset + design_resized.shape[0]
    x1, x2 = x_offset, x_offset + design_resized.shape[1]

    alpha_s = design_resized[:, :, 3] / 255.0
    alpha_l = 1.0 - alpha_s

    for c in range(0, 3):
        mockup_np[y1:y2, x1:x2, c] = (alpha_s * design_resized[:, :, c] +
                                      alpha_l * mockup_np[y1:y2, x1:x2, c])

    result_image = Image.fromarray(mockup_np)
    return result_image

# T-shirt mockup generator with Gradio interface
examples = [
    ["MyBrand"],
    ["Hello World"],
    ["Team logo"],
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # T-shirt Design Generator with Stable Diffusion
        """)

        with gr.Row():
            text = gr.Textbox(
                label="Text",
                placeholder="Enter text for the T-shirt design",
                visible=True,
            )

            run_button = gr.Button("Generate Design", scale=0)

        result = gr.Image(label="Design", show_label=False)

        gr.Examples(
            examples=examples,
            inputs=[text]
        )

    def generate_tshirt_mockup(text):
        # Generate T-shirt design
        design_image = generate_tshirt_design(text)

        # Remove background from design image
        design_np = remove_background(design_image)

        # Load blank T-shirt mockup template image
        mockup_template = Image.open("/content/drive/MyDrive/unnamed.jpg")

        # Blend design with mockup
        result_image = blend_design_with_mockup(mockup_template, Image.fromarray(design_np))

        return result_image

    run_button.click(
        fn=generate_tshirt_mockup,
        inputs=[text],
        outputs=[result]
    )

demo.queue().launch()