Qwen-Prompt-SDXL-Turbo

Runtime error

App Files Files Community

hahafofo commited on Dec 8, 2023

Commit

2b54987

1 Parent(s): 5a0f95b

init

Browse files

Files changed (3) hide show

app.py +360 -0
file_server.py +50 -0
requirements.txt +11 -0

app.py ADDED Viewed

	@@ -0,0 +1,360 @@

+# Copyright (c) Alibaba Cloud.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""A simple web interactive chat demo based on gradio."""
+import os
+from argparse import ArgumentParser
+import gradio as gr
+import mdtex2html
+import piexif
+import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.generation import GenerationConfig
+from diffusers import AutoPipelineForText2Image
+import torch
+import json
+import time
+import datetime
+from file_server import start_server, get_local_ip
+DEFAULT_CKPT_PATH = 'hahahafofo/Qwen-1_8B-Stable-Diffusion-Prompt'
+OUTPUT_IMAGES_DIR = "output_images"
+OUTPUT_HTML_DIR = "output_html"
+def _get_args():
+    parser = ArgumentParser()
+    parser.add_argument("-c", "--checkpoint-path", type=str, default=DEFAULT_CKPT_PATH,
+                        help="Checkpoint name or path, default to %(default)r")
+    parser.add_argument("--cpu-only", action="store_true", help="Run demo with CPU only")
+    parser.add_argument("--share", action="store_true", default=False,
+                        help="Create a publicly shareable link for the interface.")
+    parser.add_argument("--inbrowser", action="store_true", default=False,
+                        help="Automatically launch the interface in a new tab on the default browser.")
+    parser.add_argument("--server-port", type=int, default=8000,
+                        help="Demo server port.")
+    parser.add_argument("--server-name", type=str, default="0.0.0.0",
+                        help="Demo server name.")
+    parser.add_argument("--file-server-port", type=int, default=8001,
+                        help="file server port.")
+    args = parser.parse_args()
+    return args
+def _load_model_tokenizer(args):
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.checkpoint_path, trust_remote_code=True, resume_download=True,
+    )
+    if args.cpu_only:
+        device_map = "cpu"
+    else:
+        device_map = "auto"
+    model = AutoModelForCausalLM.from_pretrained(
+        args.checkpoint_path,
+        device_map=device_map,
+        trust_remote_code=True,
+        resume_download=True,
+    ).eval()
+    config = GenerationConfig.from_pretrained(
+        args.checkpoint_path,
+        trust_remote_code=True,
+        resume_download=True,
+    )
+    config.max_new_tokens = 100
+    return model, tokenizer, config
+def _load_sdxl_turbo():
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch.float16,
+        variant="fp16"
+    )
+    pipe.to("cuda")
+    return pipe
+def postprocess(self, y):
+    if y is None:
+        return []
+    for i, (message, response) in enumerate(y):
+        y[i] = (
+            None if message is None else mdtex2html.convert(message),
+            None if response is None else mdtex2html.convert(response),
+        )
+    return y
+gr.Chatbot.postprocess = postprocess
+def _save_image2html(image, query, prompt):
+    # 将文本信息编码为 JSON 并保存到 EXIF
+    exif_dict = {"0th": {}, "Exif": {}, "1st": {}, "thumbnail": None, "GPS": {}}
+    exif_dict["0th"][piexif.ImageIFD.ImageDescription] = json.dumps({"prompt": prompt})
+    exif_bytes = piexif.dump(exif_dict)
+    file_name = f"{int(time.time())}.png"
+    image_path = os.path.join(OUTPUT_IMAGES_DIR, file_name)
+    image.save(image_path, "PNG", exif=exif_bytes)
+    # 创建 HTML 内容
+    # 初始 HTML 结构
+    html_start = """<!DOCTYPE html><html lang="zh"><head><meta charset="UTF-8">
+    <title>Image and Prompt History</title></head><body><h1>Image and Prompt History</h1><ul>"""
+    html_end = "</ul></body></html>"
+    # 将 HTML 内容写入文件
+    html_file_path = os.path.join(OUTPUT_HTML_DIR, f"{datetime.datetime.now().strftime('%Y-%m-%d')}.html")
+    # 创建新的列表项
+    new_list_item = f"""
+        <li>
+            <p>Prompt: {prompt}</p>
+            <p>Input: {query}</p>
+            <img src="{image_path}" alt="{image_path}" style="max-width: 100%; height: auto;">
+        </li>
+    """
+    # 读取现有的 HTML 文件
+    try:
+        with open(html_file_path, 'r', encoding='utf-8') as file:
+            existing_html = file.read()
+    except FileNotFoundError:
+        # 如果文件不存在，创建一个新的 HTML 结构
+        existing_html = html_start + html_end
+    # 在列表结束标签前插入新的列表项
+    updated_html = existing_html.replace(html_end, new_list_item + html_end)
+    # 将更新后的 HTML 写回文件
+    with open(html_file_path, 'w+', encoding='utf-8') as file:
+        file.write(updated_html)
+    return f"HTML content appended to {html_file_path}"
+def _parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line != ""]
+    count = 0
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                lines[i] = f'<pre><code class="language-{items[-1]}">'
+            else:
+                lines[i] = f"<br></code></pre>"
+        else:
+            if i > 0:
+                if count % 2 == 1:
+                    line = line.replace("`", r"\`")
+                    line = line.replace("<", "&lt;")
+                    line = line.replace(">", "&gt;")
+                    line = line.replace(" ", "&nbsp;")
+                    line = line.replace("*", "&ast;")
+                    line = line.replace("_", "&lowbar;")
+                    line = line.replace("-", "&#45;")
+                    line = line.replace(".", "&#46;")
+                    line = line.replace("!", "&#33;")
+                    line = line.replace("(", "&#40;")
+                    line = line.replace(")", "&#41;")
+                    line = line.replace("$", "&#36;")
+                lines[i] = "<br>" + line
+    text = "".join(lines)
+    return text
+def _launch_demo(args, image_pipe, model, tokenizer, config):
+    def predict(
+            _query,
+            _chatbot,
+            _task_history,
+            _prompt_system: str = "You are a helpful assistant",
+            _prompt_template: str = ""
+    ):
+        print(f"User: {_parse_text(_query)}")
+        _chatbot.append((_parse_text(_query), ""))
+        full_response = ""
+        _query = f"{_prompt_template}\n{_query}"
+        for response in model.chat_stream(
+                tokenizer,
+                _query,
+                history=_task_history,
+                generation_config=config,
+                system=_prompt_system
+        ):
+            _chatbot[-1] = (_parse_text(_query), _parse_text(response))
+            yield _chatbot
+            full_response = _parse_text(response)
+        print(f"History: {_task_history}")
+        _task_history.append((_query, full_response))
+        print(f"Qwen-Chat: {_parse_text(full_response)}")
+    def draw_image(_chatbot, _task_history, num_inference_steps, ):
+        if len(_task_history) == 0:
+            return
+        prompt = _task_history[-1][-1]
+        if len(prompt) == 0:
+            return
+        print(f"===\n{_chatbot} \n\n{_task_history} ====\n")
+        print(f"{prompt}")
+        image_pil = image_pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0).images[0]
+        _save_image2html(image_pil, query=_chatbot[-1][0], prompt=prompt)
+        return image_pil
+    def regenerate(_chatbot, _task_history, _prompt_system):
+        if not _task_history:
+            yield _chatbot
+            return
+        item = _task_history.pop(-1)
+        _chatbot.pop(-1)
+        yield from predict(item[0], _chatbot, _task_history, _prompt_template="", _prompt_system=_prompt_system)
+    def reset_user_input():
+        return gr.update(value="")
+    def reset_state(_chatbot, _task_history):
+        _task_history.clear()
+        _chatbot.clear()
+        import gc
+        gc.collect()
+        torch.cuda.empty_cache()
+        return _chatbot
+    with gr.Blocks() as demo:
+        task_history = gr.State([])
+        with gr.Row():
+            with gr.Column(scale=1, min_width=600):
+                image = gr.Image(type="pil")
+                query = gr.Textbox(lines=4, label='Input')
+                with gr.Row():
+                    empty_btn = gr.Button("🧹 Clear History (清除历史)")
+                    submit_btn = gr.Button("🚀 Submit (生成)")
+                    regen_btn = gr.Button("🤔️ Regenerate (重试)")
+                    image_btn = gr.Button("🎨 Image (生成)")
+                    talk_btn = gr.Button("💬 Talk (聊天)")
+            with gr.Column(scale=1, min_width=600):
+                with gr.Tab(label="Qwen"):
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.01,
+                        value=0.9,
+                        label="Temperature",
+                        info="越小越遵循输入，越大越充满想象"
+                    )
+                    with gr.Row():
+                        prompt_system_radio = gr.Radio(
+                            ["中英文翻译", "文言文", "画家", "剧情", "AI助手"],
+                            label='角色',
+                            info="根据输入选择合适的角色"
+                        )
+                        prompt_system = gr.Textbox(
+                            lines=1,
+                            label='System Template',
+                            value="你擅长翻译中文到英语。"
+                        )
+                    prompt_template = gr.Textbox(
+                        lines=1,
+                        label='Prompt Template',
+                        value="必须使用英语根据主题描述一副画面:"
+                    )
+                    chatbot = gr.Chatbot(label='Qwen-Chat', elem_classes="control-height")
+                with gr.Tab(label="Config"):
+                    with gr.Row():
+                        top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Top-p")
+                        top_k = gr.Slider(minimum=0, maximum=100, step=1, value=50, label="Top-k")
+                        max_new_tokens = gr.Slider(minimum=1, maximum=1024, step=1, value=100, label="Max New Tokens")
+                        repetition_penalty = gr.Slider(
+                            minimum=1.0,
+                            maximum=2.0,
+                            step=0.01,
+                            value=1.1,
+                            label="repetition penalty",
+                            info="重复惩罚"
+                        )
+                    with gr.Row():
+                        num_inference_steps = gr.Slider(minimum=1, maximum=60, step=1, value=4, label="Image Steps")
+                with gr.Tab(label="History"):
+                    file_server = f"http://{get_local_ip()}:{args.file_server_port}/"
+                    html_file_path = f"{datetime.datetime.now().strftime('%Y-%m-%d')}.html"
+                    html_fns = [fn for fn in os.listdir(OUTPUT_HTML_DIR) if fn.endswith(".html")]
+                    gr.Markdown(f'<a href="{file_server}{html_file_path}" target="_blank">{html_file_path}</a>')
+                    for fn in html_fns:
+                        if fn == html_file_path:
+                            continue
+                        gr.Markdown(f'<a href="{file_server}{fn}" target="_blank">{fn}</a>')
+        PROMPT_SYSTEM_DICT = {
+            "中英文翻译": "你擅长翻译中文到英语。",
+            "文言文": "你擅长文言文翻译为英语。",
+            "画家": "你是绘画大师，擅长描绘画面细节。",
+            "剧情": "你是剧作家，擅长创作连续的漫画脚本。",
+            "AI助手": "You are a helpful assistant",
+        }
+        prompt_system_radio.change(lambda val: (PROMPT_SYSTEM_DICT[val]),
+                                   inputs=[prompt_system_radio], outputs=[prompt_system])
+        temperature.change(lambda val: config.update(temperature=val), inputs=[temperature], outputs=[])
+        top_k.change(lambda val: config.update(top_k=val), inputs=[top_k], outputs=[])
+        top_p.change(lambda val: config.update(top_p=val), inputs=[top_p], outputs=[])
+        max_new_tokens.change(
+            lambda val: config.update(max_new_tokens=val),
+            inputs=[max_new_tokens],
+            outputs=[],
+        )
+        repetition_penalty.change(
+            lambda val: config.update(repetition_penalty=val),
+            inputs=[repetition_penalty],
+            outputs=[],
+        )
+        talk_btn.click(predict, [query, chatbot, task_history, prompt_system], [chatbot],
+                       show_progress=True)
+        submit_btn.click(predict, [query, chatbot, task_history, prompt_system, prompt_template], [chatbot],
+                         show_progress=True)
+        submit_btn.click(reset_user_input, [], [query])
+        empty_btn.click(reset_state, [chatbot, task_history], outputs=[chatbot], show_progress=True)
+        image_btn.click(draw_image, [chatbot, task_history, num_inference_steps], outputs=[image],
+                        show_progress=True)
+        regen_btn.click(regenerate, [chatbot, task_history, prompt_system], [chatbot], show_progress=True)
+    demo.queue().launch(
+        share=args.share,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+        server_name=args.server_name,
+    )
+def main():
+    args = _get_args()
+    start_server(server_port=args.file_server_port)
+    os.makedirs(OUTPUT_IMAGES_DIR, exist_ok=True)
+    os.makedirs(OUTPUT_HTML_DIR, exist_ok=True)
+    model, tokenizer, config = _load_model_tokenizer(args)
+    pipe = _load_sdxl_turbo()
+    _launch_demo(args, pipe, model, tokenizer, config)
+if __name__ == '__main__':
+    main()

file_server.py ADDED Viewed

	@@ -0,0 +1,50 @@

+import http.server
+import socketserver
+import threading
+import socket
+import time
+OUTPUT_HTML_DIR = "output_html"
+def get_local_ip():
+    """ 获取当前计算机在局域网中的 IP 地址 """
+    try:
+        # 建立一个临时的连接，以便获取本地网络接口的IP地址
+        # 这里的 'www.baidu.com' 是Google的公共DNS服务器，端口为80
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        s.connect(("www.baidu.com", 80))
+        local_ip = s.getsockname()[0]
+        s.close()
+        return local_ip
+    except Exception as e:
+        print("获取本地IP地址时出错：", e)
+        return None
+class MyHttpRequestHandler(http.server.SimpleHTTPRequestHandler):
+    def do_GET(self):
+        # 指定要提供的文件的目录
+        if self.path.find(OUTPUT_HTML_DIR) == -1 and self.path.endswith(".html"):
+            self.path = OUTPUT_HTML_DIR + self.path
+        print(f"请求的文件路径为：{self.path}")
+        return http.server.SimpleHTTPRequestHandler.do_GET(self)
+def _start_server(server_port):
+    with socketserver.TCPServer(("", server_port), MyHttpRequestHandler) as httpd:
+        print(f"HTTP 文件服务器启动在端口 {server_port}")
+        httpd.serve_forever()
+def start_server(server_port):
+    # 在单独的线程中启动服务器
+    server_thread = threading.Thread(target=_start_server, args=(server_port,))
+    server_thread.daemon = True  # 设置为守护线程，这样当主程序退出时，服务器线程也会退出
+    server_thread.start()
+if __name__ == "__main__":
+    start_server(8001)
+    while True:
+        time.sleep(1000000)

requirements.txt ADDED Viewed

	@@ -0,0 +1,11 @@

+transformers==4.32.0
+accelerate
+tiktoken
+einops
+transformers_stream_generator==0.0.4
+scipy
+diffusers
+mdtex2html
+piexif
+gradio<3.42
+httpx[socks]