import os
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
import spaces

import gradio_client.utils as _gu

_orig_json_schema = _gu._json_schema_to_python_type


def _safe_json_schema(schema, defs=None):
    try:
        return _orig_json_schema(schema, defs)
    except Exception:
        return "Any"


_gu._json_schema_to_python_type = _safe_json_schema

MODEL_ID = "linhtrann21/paligemma_dual_10k"
HF_TOKEN = os.environ.get("HF_TOKEN")

print("Loading PaliGemma processor...")
processor = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224", token=HF_TOKEN)
print("Processor loaded!")

print("Loading trademark model...")
model = PaliGemmaForConditionalGeneration.from_pretrained(MODEL_ID, token=HF_TOKEN)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
print(f"Model loaded on {device}!")


@spaces.GPU
def predict(image, prompt):
    if image is None:
        return "Please upload a trademark image."

    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    image = image.convert("RGB")

    inputs = processor(
        text=prompt,
        images=image,
        return_tensors="pt"
    )

    inputs = {key: value.to(model.device) for key, value in inputs.items()}

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=5120,
            do_sample=False
        )

    input_length = inputs["input_ids"].shape[1]

    # Chỉ lấy các token được sinh thêm, bỏ phần prompt đầu vào
    generated_tokens = output[0][input_length:]

    # Chỉ decode kết quả model sinh ra
    result = processor.decode(
        generated_tokens,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True
    ).strip()

    return result


with gr.Blocks(title="Trademark Description Generator") as demo:
    gr.Markdown("# Trademark Description Generator")
    gr.Markdown("Upload a trademark image and get a detailed layout description.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Trademark Image")
            prompt_input = gr.Textbox(label="Prompt", value="Describe the layout of this trademark image.")
            submit_btn = gr.Button("Generate Description")
    
    output = gr.Textbox(label="Generated Description", lines=10)
    
    submit_btn.click(fn=predict, inputs=[image_input, prompt_input], outputs=output)


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