import gradio as gr
import tempfile
import os
import numpy as np
import torch
from PIL import Image
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation

# Проверяем наличие GPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Загружаем модель
model = TSR.from_pretrained(
    "stabilityai/TripoSR",
    config_name="config.yaml",
    weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)


# Функция для проверки изображения
def check_input_image(input_image):
    if input_image is None:
        raise gr.Error("No image uploaded!")


# Функция обработки изображения
def preprocess(input_image, do_remove_background, foreground_ratio):
    def fill_background(image):
        image = np.array(image).astype(np.float32) / 255.0
        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
        image = Image.fromarray((image * 255.0).astype(np.uint8))
        return image

    if do_remove_background:
        image = input_image.convert("RGB")
        image = remove_background(image)
        image = resize_foreground(image, foreground_ratio)
        image = fill_background(image)
    else:
        image = input_image
        if image.mode == "RGBA":
            image = fill_background(image)
    return image


# Функция генерации 3D модели
def generate(image):
    scene_codes = model(image, device=device)
    mesh = model.extract_mesh(scene_codes)[0]
    mesh = to_gradio_3d_orientation(mesh)
    mesh_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
    mesh.export(mesh_path2.name)
    return mesh_path2.name


# Настройка темы и CSS
class CustomTheme(gr.themes.Base):
    def __init__(self):
        super().__init__()
        self.primary_hue = "#191a1e"
        self.background_fill_primary = "#191a1e"
        self.background_fill_secondary = "#191a1e"
        self.background_fill_tertiary = "#191a1e"
        self.text_color_primary = "#FFFFFF"
        self.text_color_secondary = "#FFFFFF"
        self.text_color_tertiary = "#FFFFFF"
        self.input_background_fill = "#191a1e"
        self.input_text_color = "#FFFFFF"


css = """
footer {
    visibility: hidden;
    height: 0;
    margin: 0;
    padding: 0;
    overflow: hidden;
}

body, input, button, textarea, select, .gr-button {
    font-family: 'Poppins', sans-serif;
    background-color: #191a1e !important;
    color: #FFFFFF !important;
}

h1, h2, h3, h4, h5, h6 {
    font-family: 'Poppins', sans-serif;
    font-weight: 700;
    color: #FFFFFF !important;
}

input[type="text"], textarea {
    background-color: #191a1e !important;
    color: #FFFFFF !important;
    border: 1px solid #FFFFFF !important;
}

.generate-button {
    background-color: #5271FF !important;
    color: #FFFFFF !important;
    border: none !important;
    font-weight: bold !important;
}

.generate-button:hover {
    background-color: #405BBF !important;
}

gr-image-upload {
    background-color: #191a1e !important;
    border: 1px dashed #FFFFFF !important;
    position: relative;
}

.gr-image-upload .dropzone::before {
    content: 'Drop Image Here';
    color: #FFFFFF;
    font-size: 1rem;
    font-weight: bold;
    text-align: center;
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
}

.drop-image-container {
    display: flex;
    flex-direction: column;
    align-items: center;
}

.drop-image, .processed-image {
    margin-bottom: 20px;
    background-color: #191a1e !important;
}

.foreground-ratio-container {
    margin-top: 20px;
    margin-bottom: 20px;
}

.gr-tab {
    background-color: #191a1e !important;
    border: 1px solid #FFFFFF;
    border-radius: 5px;
    margin-top: 10px;
}

.gr-tab label {
    color: #FFFFFF !important;
    font-weight: bold;
}

.gr-button-primary {
    background-color: #5271FF !important;
    color: #FFFFFF !important;
    border: none;
}

.gr-button-primary:hover {
    background-color: #405BBF !important;
    color: #FFFFFF !important;
}
"""

# Интерфейс
with gr.Blocks(theme=CustomTheme(), css=css) as demo:
    with gr.Column():
        gr.Markdown("**Upload and Process Your Image**")
        with gr.Row():
            input_image = gr.Image(
                label="Upload Image",
                image_mode="RGBA",
                sources="upload",
                type="pil",
                width=400,
                height=300,
            )
            processed_image = gr.Image(
                label="Processed Image",
                interactive=False,
                width=400,
                height=300,
            )
        foreground_ratio = gr.Slider(
            label="Foreground Ratio",
            minimum=0.5,
            maximum=1.0,
            value=0.85,
            step=0.05,
        )
        do_remove_background = gr.Checkbox(
            label="Remove Background",
            value=True,
        )
        submit = gr.Button("Generate")
        output_model = gr.Model3D(label="Generated GLB Model", interactive=False)

        # Обработка кликов
        submit.click(
            fn=check_input_image,
            inputs=[input_image],
            outputs=[]
        ).then(
            fn=preprocess,
            inputs=[input_image, do_remove_background, foreground_ratio],
            outputs=[processed_image]
        ).then(
            fn=generate,
            inputs=[processed_image],
            outputs=[output_model]
        )

# Запуск приложения
demo.launch()