cppowboy/viscpm-paint

1#!/usr/bin/env python2# encoding: utf-83import gradio as gr4from datasets import load_dataset5from PIL import Image6 7import re8import os9import requests10import threading11import base64 12import io13 14sem = threading.Semaphore()15def infer(prompt, negative, scale):16    try:17        url = os.environ.get('SERVICE_URL', 'https://modelbest.cn/')18        sem.acquire()19        resp = requests.post(url, headers={20            "X-Model-Best-Model": "viscpm-paint-balance",21            "X-Model-Best-Trace-ID": "test-trace",22        }, json={23            "question": prompt,24            "negative_prompt": negative,25            "num_images_per_prompt": 4,26        })27        sem.release()28        resp = resp.json()29        images = resp['data']['response']30        images = [Image.open(io.BytesIO(base64.b64decode(encoded_image))).convert("RGB") for encoded_image in images]31        return images32    except Exception as e:33        print(e)34        return []35 36 37with gr.Blocks() as demo:38    gr.Markdown('<div align="center"><big><b>百亿参数量中英双语多模态大模型VisCPM</b></big></div>')39    with gr.Column(variant="panel"):40        with gr.Row(variant="compact"):41            with gr.Column(variant="compact"):42                text = gr.Textbox(43                    label="Enter your prompt",44                    show_label=False,45                    max_lines=1,46                    placeholder="输入提示词",47                ).style(48                    container=False,49                )50                neg_text = gr.Textbox(51                    label="Enter your negative prompt",52                    show_label=False,53                    max_lines=1,54                    placeholder="输入负向提示词（你不想生成的内容）",55                ).style(56                    container=False,57                )58            btn = gr.Button("生成图像").style(full_width=False)59 60        gallery = gr.Gallery(61            label="Generated images", show_label=False, elem_id="gallery", height="1024"62        ).style(columns=[2], rows=[2], object_fit="contain")63 64    btn.click(infer, [text, neg_text], gallery)65 66if __name__ == "__main__":67    demo.launch()