#!/usr/bin/env python
# encoding: utf-8
import gradio as gr
from datasets import load_dataset
from PIL import Image

import re
import os
import requests
import threading
import base64 
import io

sem = threading.Semaphore()
def infer(prompt, negative, scale):
    try:
        url = os.environ.get('SERVICE_URL', 'https://modelbest.cn/')
        sem.acquire()
        resp = requests.post(url, headers={
            "X-Model-Best-Model": "viscpm-paint-balance",
            "X-Model-Best-Trace-ID": "test-trace",
        }, json={
            "question": prompt,
            "negative_prompt": negative,
            "num_images_per_prompt": 4,
        })
        sem.release()
        resp = resp.json()
        images = resp['data']['response']
        images = [Image.open(io.BytesIO(base64.b64decode(encoded_image))).convert("RGB") for encoded_image in images]
        return images
    except Exception as e:
        print(e)
        return []


with gr.Blocks() as demo:
    gr.Markdown('<div align="center"><big><b>百亿参数量中英双语多模态大模型VisCPM</b></big></div>')
    with gr.Column(variant="panel"):
        with gr.Row(variant="compact"):
            with gr.Column(variant="compact"):
                text = gr.Textbox(
                    label="Enter your prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="输入提示词",
                ).style(
                    container=False,
                )
                neg_text = gr.Textbox(
                    label="Enter your negative prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="输入负向提示词（你不想生成的内容）",
                ).style(
                    container=False,
                )
            btn = gr.Button("生成图像").style(full_width=False)

        gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery", height="1024"
        ).style(columns=[2], rows=[2], object_fit="contain")

    btn.click(infer, [text, neg_text], gallery)

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