Skywork/UniPic

1import gradio as gr2from PIL import Image3import os4import tempfile5import sys6import time7from inferencer import Inferencer8from accelerate.utils import set_seed9from huggingface_hub import snapshot_download10 11model_dir = snapshot_download(repo_id="Skywork/Skywork-UniPic-1.5B")12model_path = os.path.join(model_dir,"pytorch_model.bin")13ckpt_name = "UniPic"14 15inferencer = Inferencer(16    config_file="qwen2_5_1_5b_kl16_mar_h.py",17    model_path=model_path,18    image_size=1024,19    #cfg_prompt="Generate an image.",20)21 22TEMP_DIR = tempfile.mkdtemp()23print(f"Temporary directory created at: {TEMP_DIR}")24 25def save_temp_image(pil_img):26    # 只支持512——>1024的编辑27    # img_resized = pil_img.resize((512, 512))28    path = os.path.join(TEMP_DIR, f"temp_{int(time.time())}.png")29    pil_img.save(path, format="PNG")30    return path31 32def handle_image_upload(file, history):33    if file is None:34        return None, history35    file_path = file.name if hasattr(file, "name") else file36    pil_img = Image.open(file_path)37    saved_path = save_temp_image(pil_img)38    return saved_path, history + [((saved_path,), None)]39 40def clear_all():41    for file in os.listdir(TEMP_DIR):42        path = os.path.join(TEMP_DIR, file)43        try:44            if os.path.isfile(path):45                os.remove(path)46        except Exception as e:47            print(f"Failed to delete temp file: {path}, error: {e}")48    return [], None, "Understand Image"49 50def extract_assistant_reply(full_text):51    if "assistant" in full_text:52        parts = full_text.strip().split("assistant")53        return parts[-1].lstrip(":").strip()54    return full_text.replace("<|im_end|>", "").strip()55 56def on_submit(history, user_msg, img_path, mode, grid_size=1):57    # 把 history 中的 tuples 全部换成 lists58    updated_history = [list(item) for item in history]59    user_msg = user_msg.strip()60    updated_history.append([user_msg, None])61    # set_seed(42)62 63    try:64        if mode == "Understand Image":65            if img_path is None:66                updated_history.append([None, "⚠️ Please upload or generate an image first."])67                return updated_history, "", img_path68                69            raw = (70                inferencer.query_image(Image.open(img_path), user_msg)71                if img_path else inferencer.query_text(user_msg)72            )73            reply = extract_assistant_reply(raw)74            updated_history.append([None, reply])75            return updated_history, "", img_path76 77        elif mode == "Generate Image":78            if not user_msg:79                updated_history.append([None, "⚠️ Please enter a prompt."])80                return updated_history, "", img_path81 82            imgs = inferencer.gen_image(83                raw_prompt=user_msg,84                images_to_generate=grid_size**2,85                cfg=3.0,86                num_iter=48,87                cfg_schedule="constant",88                temperature=1.0,89            )90            paths = [save_temp_image(img) for img in imgs]91            # 多图必须是列表格式92            updated_history.append([None, paths])93            return updated_history, "", paths[-1]94 95        elif mode == "Edit Image":96            if img_path is None:97                updated_history.append([None, "⚠️ Please upload or generate an image first."])98                return updated_history, "", img_path99            if not user_msg:100                updated_history.append([None, "⚠️ Please enter an edit instruction."])101                return updated_history, "", img_path102 103            img = Image.open(img_path)104            105            imgs = inferencer.edit_image(106                source_image=img,107                prompt=user_msg,108                cfg=3.0,109                cfg_prompt="repeat this image.",110                cfg_schedule="constant",111                temperature=0.85,112                grid_size=grid_size,113                num_iter=48,114            )115            paths = [save_temp_image(img) for img in imgs]116            updated_history.append([None, paths])117            return updated_history, "", paths[-1]118 119    except Exception as e:120        updated_history.append([None, f"⚠️ Failed to process: {e}"])121        return updated_history, "", img_path122 123CSS = """124/* 整体布局：上下两块 */125.gradio-container {126  display: flex !important;127  flex-direction: column;128  height: 100vh;129  margin: 0;130  padding: 0;131}132.gr-tabs {  /* ✅ 新增：确保 tab 能继承高度 */133  flex: 1 1 auto;134  display: flex;135  flex-direction: column;136}137 138/* 聊天 tab */139#tab_item_4, #tab_item_5 {140  display: flex;141  flex-direction: column;142  flex: 1 1 auto;143  overflow: hidden;  /* 防止出现双滚动条 */144  padding: 8px;145}146 147/* Chatbot 撑满 */148#chatbot1, #chatbot2{149  flex-grow: 1 !important;150  max-height: 66vh !important;  /* 限制聊天框最大高度为屏幕的2/3 */151  overflow-y: auto !important;  /* 当内容溢出时显示滚动条 */152  border: 1px solid #ddd;153  border-radius: 8px;154  padding: 12px;155  margin-bottom: 8px;156}157 158/* 图片消息放大 */159#chatbot1 img, #chatbot2 img {160  max-width: 80vw !important;161  height: auto !important;162  border-radius: 4px;163}164 165/* 底部输入区：固定高度 */166.input-row {167  flex: 0 0 auto;168  display: flex;169  align-items: center;170  padding: 8px;171  border-top: 1px solid #eee;172  background: #fafafa;173}174 175/* 文本框和按钮排布 */176.input-row .textbox-col { flex: 5; }177.input-row .upload-col, .input-row .clear-col { flex: 1; margin-left: 8px; }178 179/* 文本框样式 */180.gr-text-input {181  width: 100% !important;182  border-radius: 18px !important;183  padding: 8px 16px !important;184  border: 1px solid #ddd !important;185  font-size: 16px !important;186}187 188/* 按钮和上传组件样式 */189.gr-button, .gr-upload {190  width: 100% !important;191  border-radius: 18px !important;192  padding: 8px 16px !important;193  font-size: 16px !important;194}195"""196 197with gr.Blocks(css=CSS) as demo:198    img_state = gr.State(value=None)199    mode_state = gr.State(value="Understand Image")200 201    with gr.Tabs():202        with gr.Tab("Skywork UniPic Chatbot", elem_id="tab_item_4"):203            chatbot = gr.Chatbot(204                elem_id="chatbot1",205                show_label=False,206                avatar_images=(207                    "user.png",208                    "ai.png",209                ),210            )211            with gr.Row():212                mode_selector = gr.Radio(213                    choices=["Generate Image","Edit Image","Understand Image"],214                    value="Generate Image",215                    label="Mode",216                    interactive=True,217                )218 219            with gr.Row(elem_classes="input-row"):220                with gr.Column(elem_classes="textbox-col"):221                    user_input = gr.Textbox(222                        placeholder="Type your message here...",223                        show_label=False,224                        lines=1,225                    )226                with gr.Column(elem_classes="upload-col"):227                    image_input = gr.UploadButton(228                        "📷 Upload Image",229                        file_types=["image"],230                        file_count="single",231                        type="filepath",232                    )233                with gr.Column(elem_classes="clear-col"):234                    clear_btn = gr.Button("🧹 Clear History")235 236    user_input.submit(237        on_submit,238        [chatbot, user_input, img_state, mode_selector],239        [chatbot, user_input, img_state],240    )241    242    image_input.upload(243        handle_image_upload, [image_input, chatbot], [img_state, chatbot]244    )245    clear_btn.click(clear_all, outputs=[chatbot, img_state, mode_selector])246 247# if __name__ == "__main__":248#     demo.launch(server_name="0.0.0.0", share=True, debug=True, server_port=7689)249demo.launch()