Update app.py

app.py

@@ -1,70 +1,109 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Ovis-U1-3B  多模态 DEMO
+兼容 Hugging Face CPU Space（无 GPU 驱动）
+依赖：Python 3.10+、gradio 4.*, torch 2.*、transformers 4.41.*
+"""
+
 import os
 import subprocess
-subprocess.run('pip install flash-attn==2.6.3 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 import random
-import spaces
 import numpy as np
 import torch
 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
-
-
-# Constants
-MAX_SEED = 10000
 
-hf_token = os.getenv("HF_TOKEN")
-
-HUB_MODEL_ID = "AIDC-AI/Ovis-U1-3B"
-model, loading_info = AutoModelForCausalLM.from_pretrained(
-    HUB_MODEL_ID, 
-    torch_dtype=torch.bfloat16,
-    output_loading_info=True,
-    token=hf_token,
-    trust_remote_code=True
+# -------------------------------------------------------------------------
+# 可选：GPU 环境才能真正用到 flash-attn；CPU Space 可忽略安装异常
+# -------------------------------------------------------------------------
+try:
+    subprocess.run(
+        "pip install flash-attn==2.6.3 --no-build-isolation",
+        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+        shell=True,
+        check=True,
     )
-print(f'Loading info of Ovis-U1:\n{loading_info}')
-
-model = model.eval().to("cuda")
-model = model.to(torch.bfloat16)
-
-def set_global_seed(seed: int = 42):
+except subprocess.CalledProcessError:
+    print("[INFO] flash-attn 安装失败（CPU 环境可忽略）")
+
+# -------------------------------------------------------------------------
+# 常量与工具函数
+# -------------------------------------------------------------------------
+MAX_SEED = 10_000
+DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE    = torch.bfloat16 if DEVICE == "cuda" else torch.float32   # CPU 默认用 fp32
+
+def set_global_seed(seed: int = 42) -> None:
+    """统一设置随机种子（CPU / CUDA 自适应）"""
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
 def randomize_seed_fn(seed: int, randomize: bool) -> int:
+    """UI 侧 seed 随机化"""
     return random.randint(0, MAX_SEED) if randomize else seed
 
-@spaces.GPU
-def process_txt_to_img(prompt: str, height: int, width: int, steps: int, final_seed: int, guidance_scale: float, progress: gr.Progress = gr.Progress(track_tqdm=True)) -> list[Image.Image]:
+# -------------------------------------------------------------------------
+# 加载模型
+# -------------------------------------------------------------------------
+HF_TOKEN      = os.getenv("HF_TOKEN")          # 如果私有模型需 token
+HUB_MODEL_ID  = "AIDC-AI/Ovis-U1-3B"
+
+print(f"[INFO] Loading {HUB_MODEL_ID} on {DEVICE} ...")
+model = AutoModelForCausalLM.from_pretrained(
+    HUB_MODEL_ID,
+    torch_dtype=DTYPE,
+    low_cpu_mem_usage=True,    # 显著降低 CPU 占用
+    device_map="auto",         # cuda 自动放 GPU，CPU 环境全部放 CPU
+    token=HF_TOKEN,
+    trust_remote_code=True
+).eval()                       # 评估模式
+
+print("[INFO] Model ready!")
+
+# -------------------------------------------------------------------------
+# 引入自定义管线函数 —— 保持与原代码一致
+# -------------------------------------------------------------------------
+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
+
+# -------------------------------------------------------------------------
+# 推理封装（均运行在 DEVICE）
+# -------------------------------------------------------------------------
+def process_txt_to_img(prompt: str, height: int, width: int, steps: int,
+                       final_seed: int, guidance_scale: float,
+                       progress: gr.Progress = gr.Progress(track_tqdm=True)) -> list[Image.Image]:
     set_global_seed(final_seed)
-    images = pipe_t2i(model, prompt, height, width, steps, cfg=guidance_scale, seed=final_seed)
-    return images
+    return pipe_t2i(model, prompt, height, width, steps,
+                    cfg=guidance_scale, seed=final_seed)
 
-@spaces.GPU
-def process_img_to_txt(prompt: str, img: Image.Image, progress: gr.Progress = gr.Progress(track_tqdm=True)) -> str:
-    output_text = pipe_txt_gen(model, img, prompt)
-    return output_text
+def process_img_to_txt(prompt: str, img: Image.Image,
+                       progress: gr.Progress = gr.Progress(track_tqdm=True)) -> str:
+    return pipe_txt_gen(model, img, prompt)
 
-@spaces.GPU
-def process_img_txt_to_img(prompt: str, img: Image.Image, steps: int, final_seed: int, txt_cfg: float, img_cfg: float, progress: gr.Progress = gr.Progress(track_tqdm=True)) -> list[Image.Image]:
+def process_img_txt_to_img(prompt: str, img: Image.Image, steps: int,
+                           final_seed: int, txt_cfg: float, img_cfg: float,
+                           progress: gr.Progress = gr.Progress(track_tqdm=True)) -> list[Image.Image]:
     set_global_seed(final_seed)
-    images = pipe_img_edit(model, img, prompt, steps, txt_cfg, img_cfg, seed=final_seed)
-    return images
+    return pipe_img_edit(model, img, prompt, steps,
+                         txt_cfg, img_cfg, seed=final_seed)
 
+# -------------------------------------------------------------------------
 # Gradio UI
-with gr.Blocks(title="Ovis-U1-3B") as demo:
-    gr.Markdown('''# Ovis-U1-3B
-    ''')
+# -------------------------------------------------------------------------
+with gr.Blocks(title="Ovis-U1-3B (CPU)") as demo:
+    gr.Markdown("# Ovis-U1-3B\n✨ 多模态文本-图像 DEMO（CPU 版）")
 
     with gr.Row():
         with gr.Column():
             with gr.Tabs():
+                # ---------------------- Tab 1  图像 + 文本 → 图像 ----------------------
                 with gr.TabItem("Image + Text → Image"):
                     edit_image_input = gr.Image(label="Input Image", type="pil")
                     with gr.Row():
@@ -78,47 +117,38 @@ with gr.Blocks(title="Ovis-U1-3B") as demo:
                         run_edit_image_btn = gr.Button("Run", scale=0)
 
                     with gr.Accordion("Advanced Settings", open=False):
-
                         with gr.Row():
-
                             edit_img_guidance_slider = gr.Slider(
                                 label="Image Guidance Scale",
-                                minimum=1.0, maximum=10.0,
-                                step=0.1, value=1.5
+                                minimum=1.0, maximum=10.0, step=0.1, value=1.5
                             )
-
                             edit_txt_guidance_slider = gr.Slider(
                                 label="Text Guidance Scale",
-                                minimum=1.0, maximum=30.0,
-                                step=0.5, value=6.0
+                                minimum=1.0, maximum=30.0, step=0.5, value=6.0
                             )
-
                         edit_num_steps_slider = gr.Slider(
-                            label='Steps',
-                            minimum=40, maximum=100, 
-                            value=50, step=1
+                            label="Steps", minimum=40, maximum=100, value=50, step=1
                         )
                         edit_seed_slider = gr.Slider(
-                            label="Seed",
-                            minimum=0, maximum=int(MAX_SEED),
+                            label="Seed", minimum=0, maximum=MAX_SEED,
                             step=1, value=42
                         )
                         edit_randomize_checkbox = gr.Checkbox(
                             label="Randomize seed", value=False
                         )
 
-                    img_edit_examples_data = [
-                        ["imgs/train.png", "Modify this image in a Ghibli style. "],
-                        ["imgs/chair.png", "Transfer the image into a faceted low-poly 3-D render style."],
-                        ["imgs/car.png", "Replace the tiny house on wheels in the image with a vintage car."],
-                    ]
                     gr.Examples(
-                        examples=img_edit_examples_data,
+                        examples=[
+                            ["imgs/train.png", "Modify this image in a Ghibli style."],
+                            ["imgs/chair.png", "Transfer the image into a faceted low-poly 3-D render style."],
+                            ["imgs/car.png", "Replace the tiny house on wheels in the image with a vintage car."],
+                        ],
                         inputs=[edit_image_input, edit_prompt_input],
-                        cache_examples=False, 
+                        cache_examples=False,
                         label="Image Editing Examples"
                     )
 
+                # ---------------------- Tab 2  文本 → 图像 ----------------------
                 with gr.TabItem("Text → Image"):
                     with gr.Row():
                         prompt_gen_input = gr.Textbox(
@@ -133,212 +163,188 @@ with gr.Blocks(title="Ovis-U1-3B") as demo:
                     with gr.Accordion("Advanced Settings", open=False):
                         with gr.Row():
                             height_slider = gr.Slider(
-                                label='height', 
-                                minimum=256, maximum=1536, 
+                                label="height", minimum=256, maximum=1536,
                                 value=1024, step=32
                             )
                             width_slider = gr.Slider(
-                                label='width', 
-                                minimum=256, maximum=1536, 
+                                label="width", minimum=256, maximum=1536,
                                 value=1024, step=32
                             )
-
                         guidance_slider = gr.Slider(
                             label="Guidance Scale",
-                            minimum=1.0, maximum=30.0,
-                            step=0.5, value=5.0
+                            minimum=1.0, maximum=30.0, step=0.5, value=5.0
                         )
-                        
                         num_steps_slider = gr.Slider(
-                            label='Steps',
-                            minimum=40, maximum=100, 
-                            value=50, step=1
+                            label="Steps", minimum=40, maximum=100, value=50, step=1
                         )
                         seed_slider = gr.Slider(
-                            label="Seed",
-                            minimum=0, maximum=int(MAX_SEED),
+                            label="Seed", minimum=0, maximum=MAX_SEED,
                             step=1, value=42
                         )
                         randomize_checkbox = gr.Checkbox(
                             label="Randomize seed", value=False
                         )
 
-                    text_gen_examples_data = [
-                        ["A breathtaking fairy with teal wings sits gracefully on a lotus flower in a serene pond, exuding elegance."],
-                        ["A winter mountain landscape at deep night with snowy terrain and colorful flowers, under beautiful clouds and no people, portrayed as an anime background illustration with intricate detail and sharp focus."],
-                        ["A photo of a pug wearing a cowboy hat and bandana, sitting on a hay bale."]
-                    ]
                     gr.Examples(
-                        examples=text_gen_examples_data,
+                        examples=[
+                            ["A breathtaking fairy with teal wings sits gracefully on a lotus flower in a serene pond, exuding elegance."],
+                            ["A winter mountain landscape at deep night with snowy terrain and colorful flowers, portrayed as an anime background illustration."],
+                            ["A photo of a pug wearing a cowboy hat and bandana, sitting on a hay bale."]
+                        ],
                         inputs=[prompt_gen_input],
-                        cache_examples=False, 
+                        cache_examples=False,
                         label="Image Generation Examples"
                     )
 
+                # ---------------------- Tab 3  图像 → 文本 ----------------------
                 with gr.TabItem("Image → Text"):
                     image_understand_input = gr.Image(label="Input Image", type="pil")
                     with gr.Row():
                         prompt_understand_input = gr.Textbox(
-                            label="Prompt",
-                            show_label=False,
+                            label="Prompt", show_label=False,
                             placeholder="Describe the question about image...",
-                            container=False,
-                            lines=1
+                            container=False, lines=1
                         )
                         run_image_understand_btn = gr.Button("Run", scale=0)
 
-                    image_understanding_examples_data = [
-                        ["imgs/table.webp", "In what scenario does this picture take place?"],
-                        ["imgs/count.png", "How many broccoli are there in the picture?"],
-                        ["imgs/foot.webp", "Where is this picture located?"],
-                    ]
                     gr.Examples(
-                        examples=image_understanding_examples_data,
+                        examples=[
+                            ["imgs/table.webp", "In what scenario does this picture take place?"],
+                            ["imgs/count.png", "How many broccoli are there in the picture?"],
+                            ["imgs/foot.webp", "Where is this picture located?"],
+                        ],
                         inputs=[image_understand_input, prompt_understand_input],
-                        cache_examples=False, 
+                        cache_examples=False,
                         label="Image Understanding Examples"
                     )
-            
-            clean_btn  = gr.Button("Clear All Inputs/Outputs")
 
-        with gr.Column():
-            output_gallery = gr.Gallery(label="Generated Images", columns=2, visible=True) # Default to visible, content will control
-            output_text    = gr.Textbox(label="Generated Text", visible=False, lines=5, interactive=False)
+            clean_btn = gr.Button("Clear All Inputs / Outputs")
 
-    @spaces.GPU
-    def run_img_txt_to_img_tab(prompt, img, steps, seed, txt_cfg, img_cfg, progress=gr.Progress(track_tqdm=True)):
+        # ---------------------- 输出区 ----------------------
+        with gr.Column():
+            output_gallery = gr.Gallery(label="Generated Images",
+                                        columns=2, visible=True)
+            output_text = gr.Textbox(label="Generated Text",
+                                     visible=False, lines=5,
+                                     interactive=False)
+
+    # ---------------------------------------------------------------------
+    # 事件绑定
+    # ---------------------------------------------------------------------
+    def run_img_txt_to_img_tab(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 for editing.", visible=True)
             )
-        # Seed is already finalized by the randomize_seed_fn in the click chain
-        imgs = process_img_txt_to_img(prompt, img, steps, seed, txt_cfg, img_cfg, progress=progress)
+        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)
         )
 
-    @spaces.GPU
-    def run_txt_to_img_tab(prompt, height, width, steps, seed, guidance, progress=gr.Progress(track_tqdm=True)):
-        # Seed is already finalized by the randomize_seed_fn in the click chain
-        imgs = process_txt_to_img(prompt, height, width, steps, seed, guidance, progress=progress)
+    def run_txt_to_img_tab(prompt, height, width, steps, seed, guidance,
+                           progress=gr.Progress(track_tqdm=True)):
+        imgs = process_txt_to_img(prompt, height, width, steps, seed,
+                                  guidance, progress)
         return (
             gr.update(value=imgs, visible=True),
             gr.update(value="", visible=False)
         )
 
-    @spaces.GPU
-    def run_img_to_txt_tab(img, prompt, progress=gr.Progress(track_tqdm=True)):
+    def run_img_to_txt_tab(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 for understanding.", visible=True)
             )
-        txt = process_img_to_txt(prompt, img, progress=progress)
+        txt = process_img_to_txt(prompt, img, progress)
         return (
             gr.update(value=[], visible=False),
             gr.update(value=txt, visible=True)
         )
 
     def clean_all_fn():
+        """重置全部输入 / 输出"""
         return (
-            # Tab 1 inputs
-            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),
+            # Tab 1
+            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),
-            # Tab 2 inputs
-            gr.update(value=""),  # prompt_gen_input
-            gr.update(value=1024),
-            gr.update(value=1024),
-            gr.update(value=5.0),
-            gr.update(value=50),
-            gr.update(value=42),  # seed_slider
-            gr.update(value=False), # randomize_checkbox
-            # Tab 3 inputs
-            gr.update(value=None), # image_understand_input
-            gr.update(value=""),  # prompt_understand_input
+            # Tab 2
+            gr.update(value=""),   gr.update(value=1024),
+            gr.update(value=1024), gr.update(value=5.0),
+            gr.update(value=50),   gr.update(value=42),
+            gr.update(value=False),
+            # Tab 3
+            gr.update(value=None), gr.update(value=""),
             # Outputs
-            gr.update(value=[], visible=True), # output_gallery (reset and keep visible for next gen)
-            gr.update(value="", visible=False) # output_text (reset and hide)
+            gr.update(value=[], visible=True),
+            gr.update(value="",  visible=False)
         )
 
-    # Event listeners for Image + Text -> Image
-    edit_inputs = [edit_prompt_input, edit_image_input, edit_num_steps_slider, edit_seed_slider, edit_txt_guidance_slider, edit_img_guidance_slider]
-    
-    run_edit_image_btn.click(
-        fn=randomize_seed_fn,
-        inputs=[edit_seed_slider, edit_randomize_checkbox],
-        outputs=[edit_seed_slider]
-    ).then(
-        fn=run_img_txt_to_img_tab,
-        inputs=edit_inputs, 
-        outputs=[output_gallery, output_text]
+    # ---------- Tab 1 ----------
+    edit_inputs = [
+        edit_prompt_input, edit_image_input,
+        edit_num_steps_slider, edit_seed_slider,
+        edit_txt_guidance_slider, edit_img_guidance_slider
+    ]
+    run_edit_image_btn.click(randomize_seed_fn,
+                             [edit_seed_slider, edit_randomize_checkbox],
+                             [edit_seed_slider]).then(
+        run_img_txt_to_img_tab, edit_inputs,
+        [output_gallery, output_text]
     )
-
-    edit_prompt_input.submit(
-        fn=randomize_seed_fn,
-        inputs=[edit_seed_slider, edit_randomize_checkbox],
-        outputs=[edit_seed_slider]
-    ).then(
-        fn=run_img_txt_to_img_tab,
-        inputs=edit_inputs,
-        outputs=[output_gallery, output_text]
+    edit_prompt_input.submit(randomize_seed_fn,
+                             [edit_seed_slider, edit_randomize_checkbox],
+                             [edit_seed_slider]).then(
+        run_img_txt_to_img_tab, edit_inputs,
+        [output_gallery, output_text]
     )
 
-    # Event listeners for Text -> Image
-    gen_inputs = [prompt_gen_input, height_slider, width_slider, num_steps_slider, seed_slider, guidance_slider]
-    
-    run_image_gen_btn.click(
-        fn=randomize_seed_fn,
-        inputs=[seed_slider, randomize_checkbox],
-        outputs=[seed_slider]
-    ).then(
-        fn=run_txt_to_img_tab,
-        inputs=gen_inputs, 
-        outputs=[output_gallery, output_text]
+    # ---------- Tab 2 ----------
+    gen_inputs = [
+        prompt_gen_input, height_slider, width_slider,
+        num_steps_slider, seed_slider, guidance_slider
+    ]
+    run_image_gen_btn.click(randomize_seed_fn,
+                            [seed_slider, randomize_checkbox],
+                            [seed_slider]).then(
+        run_txt_to_img_tab, gen_inputs,
+        [output_gallery, output_text]
     )
-
-    prompt_gen_input.submit(
-        fn=randomize_seed_fn,
-        inputs=[seed_slider, randomize_checkbox],
-        outputs=[seed_slider]
-    ).then(
-        fn=run_txt_to_img_tab,
-        inputs=gen_inputs,
-        outputs=[output_gallery, output_text]
+    prompt_gen_input.submit(randomize_seed_fn,
+                            [seed_slider, randomize_checkbox],
+                            [seed_slider]).then(
+        run_txt_to_img_tab, gen_inputs,
+        [output_gallery, output_text]
     )
 
-    # Event listeners for Image -> Text
+    # ---------- Tab 3 ----------
     understand_inputs = [image_understand_input, prompt_understand_input]
-
-    run_image_understand_btn.click(
-        fn=run_img_to_txt_tab,
-        inputs=understand_inputs,
-        outputs=[output_gallery, output_text]
-    )
-
-    prompt_understand_input.submit(
-        fn=run_img_to_txt_tab,
-        inputs=understand_inputs,
-        outputs=[output_gallery, output_text]
-    )
-
-    clean_btn.click(
-        fn=clean_all_fn,
-        inputs=[],
-        outputs=[
-            edit_image_input, edit_prompt_input, edit_img_guidance_slider, edit_txt_guidance_slider,
-            edit_num_steps_slider, edit_seed_slider, edit_randomize_checkbox,
-            prompt_gen_input, height_slider, width_slider, guidance_slider, num_steps_slider, seed_slider, randomize_checkbox, 
-            image_understand_input, prompt_understand_input,
-            output_gallery, output_text
-        ]
-    )
-
+    run_image_understand_btn.click(run_img_to_txt_tab,
+                                   understand_inputs,
+                                   [output_gallery, output_text])
+    prompt_understand_input.submit(run_img_to_txt_tab,
+                                   understand_inputs,
+                                   [output_gallery, output_text])
+
+    # ---------- 清空 ----------
+    clean_btn.click(clean_all_fn, [], [
+        edit_image_input, edit_prompt_input, edit_img_guidance_slider,
+        edit_txt_guidance_slider, edit_num_steps_slider, edit_seed_slider,
+        edit_randomize_checkbox, prompt_gen_input, height_slider,
+        width_slider, guidance_slider, num_steps_slider, seed_slider,
+        randomize_checkbox, image_understand_input, prompt_understand_input,
+        output_gallery, output_text
+    ])
+
+# -------------------------------------------------------------------------
+# 启动
+# -------------------------------------------------------------------------
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)     # HF Spaces 自动监听 7860 端口