#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Ovis-U1-3B 多模态 DEMO（CPU / GPU 自适应版本） 依赖：Python 3.10+、torch 2.*、transformers 4.41.*、gradio 4.* """ # ─────────────────────────────────────────────── # ① 在任何 transformers / flash_attn 导入之前处理环境 # ─────────────────────────────────────────────── import os, sys, types, subprocess, random, numpy as np, torch import importlib.util # ★ 新增：用于生成 ModuleSpec DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32 # -------- CPU 环境：屏蔽 flash-attn -------- if DEVICE == "cpu": # 卸载潜在的 flash-attn subprocess.run("pip uninstall -y flash-attn", shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # 构造空壳模块 fake_flash_attn = types.ModuleType("flash_attn") fake_layers = types.ModuleType("flash_attn.layers") fake_rotary = types.ModuleType("flash_attn.layers.rotary") def _cpu_apply_rotary_emb(x, cos, sin): """纯 CPU 的旋转位置编码（简易实现）""" x1, x2 = x[..., ::2], x[..., 1::2] rot_x1 = x1 * cos - x2 * sin rot_x2 = x1 * sin + x2 * cos out = torch.empty_like(x) out[..., ::2] = rot_x1 out[..., 1::2] = rot_x2 return out fake_rotary.apply_rotary_emb = _cpu_apply_rotary_emb fake_layers.rotary = fake_rotary fake_flash_attn.layers = fake_layers # ★ 新增：为空壳模块补充合法的 __spec__ fake_flash_attn.__spec__ = importlib.util.spec_from_loader("flash_attn", loader=None) sys.modules.update({ "flash_attn": fake_flash_attn, "flash_attn.layers": fake_layers, "flash_attn.layers.rotary": fake_rotary, }) else: # GPU 环境：尝试安装 flash-attn try: subprocess.run( "pip install flash-attn==2.6.3 --no-build-isolation", env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}, shell=True, check=True) except subprocess.CalledProcessError: print("[WARN] flash-attn 安装失败，GPU 加速功能受限。") # ─────────────────────────────────────────────── # ② 常规依赖 # ─────────────────────────────────────────────── from PIL import Image import gradio as gr import spaces from transformers import AutoModelForCausalLM from test_img_edit import pipe_img_edit from test_img_to_txt import pipe_txt_gen from test_txt_to_img import pipe_t2i # ─────────────────────────────────────────────── # ③ 工具函数 & 常量 # ─────────────────────────────────────────────── MAX_SEED = 10_000 def set_global_seed(seed: int = 42): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def randomize_seed_fn(seed: int, randomize: bool) -> int: return random.randint(0, MAX_SEED) if randomize else seed # ─────────────────────────────────────────────── # ④ 加载模型 # ─────────────────────────────────────────────── HF_TOKEN = os.getenv("HF_TOKEN") MODEL_ID = "AIDC-AI/Ovis-U1-3B" print(f"[INFO] Loading {MODEL_ID} on {DEVICE} …") model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=DTYPE, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN, trust_remote_code=True ).eval() print("[INFO] Model ready!") # ─────────────────────────────────────────────── # ⑤ 推理封装 # ─────────────────────────────────────────────── def process_txt_to_img(prompt, height, width, steps, seed, cfg, progress=gr.Progress(track_tqdm=True)): set_global_seed(seed) return pipe_t2i(model, prompt, height, width, steps, cfg=cfg, seed=seed) def process_img_to_txt(prompt, img, progress=gr.Progress(track_tqdm=True)): return pipe_txt_gen(model, img, prompt) def process_img_txt_to_img(prompt, img, steps, seed, txt_cfg, img_cfg, progress=gr.Progress(track_tqdm=True)): set_global_seed(seed) return pipe_img_edit(model, img, prompt, steps, txt_cfg, img_cfg, seed=seed) # ─────────────────────────────────────────────── # ⑥ Gradio UI（与前版一致，此处省略修改标记） # ─────────────────────────────────────────────── with gr.Blocks(title="Ovis-U1-3B (CPU/GPU adaptive)") as demo: gr.Markdown("# Ovis-U1-3B\n多模态文本-图像 DEMO（CPU/GPU 自适应版）") with gr.Row(): with gr.Column(): with gr.Tabs(): # Tab 1: Image + Text → Image with gr.TabItem("Image + Text → Image"): edit_image_input = gr.Image(label="Input Image", type="pil") with gr.Row(): edit_prompt_input = gr.Textbox(show_label=False, placeholder="Describe the editing instruction…") run_edit_image_btn = gr.Button("Run", scale=0) with gr.Accordion("Advanced Settings", open=False): with gr.Row(): edit_img_guidance = gr.Slider(label="Image Guidance", minimum=1, maximum=10, value=1.5, step=0.1) edit_txt_guidance = gr.Slider(label="Text Guidance", minimum=1, maximum=30, value=6.0, step=0.5) edit_steps = gr.Slider(label="Steps", minimum=40, maximum=100, value=50, step=1) edit_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1) edit_random = gr.Checkbox(label="Randomize seed", value=False) # Tab 2: Text → Image with gr.TabItem("Text → Image"): prompt_gen = gr.Textbox(show_label=False, placeholder="Describe the image you want…") run_gen_btn = gr.Button("Run", scale=0) with gr.Accordion("Advanced Settings", open=False): with gr.Row(): height_slider = gr.Slider(label="height", minimum=256, maximum=1536, value=1024, step=32) width_slider = gr.Slider(label="width", minimum=256, maximum=1536, value=1024, step=32) guidance_slider = gr.Slider(label="Guidance Scale", minimum=1, maximum=30, value=5, step=0.5) steps_slider = gr.Slider(label="Steps", minimum=40, maximum=100, value=50, step=1) seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1) random_check = gr.Checkbox(label="Randomize seed", value=False) # Tab 3: Image → Text with gr.TabItem("Image → Text"): understand_img = gr.Image(label="Input Image", type="pil") understand_prompt = gr.Textbox(show_label=False, placeholder="Describe the question about image…") run_understand = gr.Button("Run", scale=0) clear_btn = gr.Button("Clear All") with gr.Column(): gallery = gr.Gallery(label="Generated Images", columns=2, visible=True) txt_out = gr.Textbox(label="Generated Text", visible=False, lines=5, interactive=False) # 事件绑定（与上一版相同，省略重复注释） def run_tab1(prompt, img, steps, seed, txt_cfg, img_cfg, progress=gr.Progress(track_tqdm=True)): if img is None: return gr.update(value=[], visible=False), gr.update(value="Please upload an image.", visible=True) imgs = process_img_txt_to_img(prompt, img, steps, seed, txt_cfg, img_cfg, progress) return gr.update(value=imgs, visible=True), gr.update(value="", visible=False) def run_tab2(prompt, h, w, steps, seed, guidance, progress=gr.Progress(track_tqdm=True)): imgs = process_txt_to_img(prompt, h, w, steps, seed, guidance, progress) return gr.update(value=imgs, visible=True), gr.update(value="", visible=False) def run_tab3(img, prompt, progress=gr.Progress(track_tqdm=True)): if img is None: return gr.update(value=[], visible=False), gr.update(value="Please upload an image.", visible=True) text = process_img_to_txt(prompt, img, progress) return gr.update(value=[], visible=False), gr.update(value=text, visible=True) # Tab1 绑定 run_edit_image_btn.click(randomize_seed_fn, [edit_seed, edit_random], [edit_seed]).then( run_tab1, [edit_prompt_input, edit_image_input, edit_steps, edit_seed, edit_txt_guidance, edit_img_guidance], [gallery, txt_out]) edit_prompt_input.submit(randomize_seed_fn, [edit_seed, edit_random], [edit_seed]).then( run_tab1, [edit_prompt_input, edit_image_input, edit_steps, edit_seed, edit_txt_guidance, edit_img_guidance], [gallery, txt_out]) # Tab2 绑定 run_gen_btn.click(randomize_seed_fn, [seed_slider, random_check], [seed_slider]).then( run_tab2, [prompt_gen, height_slider, width_slider, steps_slider, seed_slider, guidance_slider], [gallery, txt_out]) prompt_gen.submit(randomize_seed_fn, [seed_slider, random_check], [seed_slider]).then( run_tab2, [prompt_gen, height_slider, width_slider, steps_slider, seed_slider, guidance_slider], [gallery, txt_out]) # Tab3 绑定 run_understand.click(run_tab3, [understand_img, understand_prompt], [gallery, txt_out]) understand_prompt.submit(run_tab3, [understand_img, understand_prompt], [gallery, txt_out]) # 清空 def clear_all(): return ( gr.update(value=None), gr.update(value=""), gr.update(value=1.5), gr.update(value=6.0), gr.update(value=50), gr.update(value=42), gr.update(value=False), gr.update(value=""), gr.update(value=1024), gr.update(value=1024), gr.update(value=5), gr.update(value=50), gr.update(value=42), gr.update(value=False), gr.update(value=None), gr.update(value=""), gr.update(value=[], visible=True), gr.update(value="", visible=False) ) clear_btn.click(clear_all, [], [ edit_image_input, edit_prompt_input, edit_img_guidance, edit_txt_guidance, edit_steps, edit_seed, edit_random, prompt_gen, height_slider, width_slider, guidance_slider, steps_slider, seed_slider, random_check, understand_img, understand_prompt, gallery, txt_out ]) # ─────────────────────────────────────────────── # ⑦ 启动 # ─────────────────────────────────────────────── if __name__ == "__main__": demo.launch()