Upload 12 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+imgs/car.png filter=lfs diff=lfs merge=lfs -text
+imgs/chair.png filter=lfs diff=lfs merge=lfs -text
+imgs/count.png filter=lfs diff=lfs merge=lfs -text
+imgs/foot.webp filter=lfs diff=lfs merge=lfs -text
+imgs/table.webp filter=lfs diff=lfs merge=lfs -text
+imgs/train.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,14 @@
 ---
 title: Ovis U1 3B
-emoji: 🦀
-colorFrom: gray
-colorTo: pink
+emoji: 🎨
+colorFrom: green
+colorTo: indigo
 sdk: gradio
 sdk_version: 5.35.0
 app_file: app.py
 pinned: false
 license: apache-2.0
+short_description: Demo for multimodal understanding and generation
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,344 @@
+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
+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
+    )
+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):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    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
+
+@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]:
+    set_global_seed(final_seed)
+    images = pipe_t2i(model, prompt, height, width, steps, cfg=guidance_scale, seed=final_seed)
+    return images
+
+@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
+
+@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]:
+    set_global_seed(final_seed)
+    images = pipe_img_edit(model, img, prompt, steps, txt_cfg, img_cfg, seed=final_seed)
+    return images
+
+# Gradio UI
+with gr.Blocks(title="Ovis-U1-3B") as demo:
+    gr.Markdown('''# Ovis-U1-3B
+    ''')
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs():
+                with gr.TabItem("Image + Text → Image"):
+                    edit_image_input = gr.Image(label="Input Image", type="pil")
+                    with gr.Row():
+                        edit_prompt_input = gr.Textbox(
+                            label="Prompt",
+                            show_label=False,
+                            placeholder="Describe the editing instruction...",
+                            container=False,
+                            lines=1
+                        )
+                        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
+                            )
+
+                            edit_txt_guidance_slider = gr.Slider(
+                                label="Text Guidance Scale",
+                                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
+                        )
+                        edit_seed_slider = gr.Slider(
+                            label="Seed",
+                            minimum=0, maximum=int(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,
+                        inputs=[edit_image_input, edit_prompt_input],
+                        cache_examples=False, 
+                        label="Image Editing Examples"
+                    )
+
+                with gr.TabItem("Text → Image"):
+                    with gr.Row():
+                        prompt_gen_input = gr.Textbox(
+                            label="Prompt",
+                            show_label=False,
+                            placeholder="Describe the image you want...",
+                            container=False,
+                            lines=1
+                        )
+                        run_image_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.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
+                        )
+                        seed_slider = gr.Slider(
+                            label="Seed",
+                            minimum=0, maximum=int(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,
+                        inputs=[prompt_gen_input],
+                        cache_examples=False, 
+                        label="Image Generation Examples"
+                    )
+
+                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,
+                            placeholder="Describe the question about image...",
+                            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,
+                        inputs=[image_understand_input, prompt_understand_input],
+                        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)
+
+    @spaces.GPU
+    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)
+        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)
+        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)):
+        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)
+        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),
+            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
+            # 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)
+        )
+
+    # 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]
+    )
+
+    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]
+    )
+
+    # 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]
+    )
+
+    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]
+    )
+
+    # Event listeners for Image -> Text
+    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
+        ]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,29 @@
+torch==2.4.0
+transformers==4.51.3
+tokenizers==0.21.1
+sentencepiece==0.1.99
+pyarrow==18.0.0
+accelerate==1.1.0
+pydantic==2.8.2
+markdown2[all]
+numpy==1.24.3
+scikit-learn==1.2.2
+requests
+httpx
+uvicorn
+fastapi==0.112.4
+einops==0.6.1
+einops-exts==0.0.4
+timm==1.0.11
+tiktoken
+transformers_stream_generator==0.0.4
+scipy
+pandas
+torchaudio
+xformers
+pillow==10.3.0
+pysubs2==1.7.2
+trl==0.12.1
+moviepy==1.0.3
+diffusers==0.31.0
+gradio

test_img_edit.py

@@ -0,0 +1,132 @@
+import os
+import argparse
+import math
+import numpy as np
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Test Image Editing")
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="AIDC-AI/Ovis-U1-3B",
+    )
+    parser.add_argument(
+        "--steps", type=int, default=50,
+    )
+    parser.add_argument(
+        "--img_cfg", type=float, default=1.5,
+    )
+    parser.add_argument(
+        "--txt_cfg", type=float, default=6,
+    )
+    args = parser.parse_args()
+    return args
+
+def load_blank_image(width, height):
+    pil_image = Image.new("RGB", (width, height), (255, 255, 255)).convert('RGB')
+    return pil_image
+
+def build_inputs(model, text_tokenizer, visual_tokenizer, prompt, pil_image, target_width, target_height):
+    if pil_image is not None:
+        target_size = (int(target_width), int(target_height))
+        pil_image, vae_pixel_values, cond_img_ids = model.visual_generator.process_image_aspectratio(pil_image, target_size)
+        cond_img_ids[..., 0] = 1.0
+        vae_pixel_values = vae_pixel_values.unsqueeze(0).to(device=model.device)
+        width = pil_image.width
+        height = pil_image.height
+        resized_height, resized_width = visual_tokenizer.smart_resize(height, width, max_pixels=visual_tokenizer.image_processor.min_pixels)
+        pil_image = pil_image.resize((resized_width, resized_height))
+    else:
+        vae_pixel_values = None
+        cond_img_ids = None
+
+    prompt, input_ids, pixel_values, grid_thws = model.preprocess_inputs(
+        prompt, 
+        [pil_image], 
+        generation_preface=None,
+        return_labels=False,
+        propagate_exception=False,
+        multimodal_type='single_image',
+        fix_sample_overall_length_navit=False
+        )
+    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
+    input_ids = input_ids.unsqueeze(0).to(device=model.device)
+    attention_mask = attention_mask.unsqueeze(0).to(device=model.device)
+    if pixel_values is not None:
+        pixel_values = torch.cat([
+            pixel_values.to(device=visual_tokenizer.device, dtype=torch.bfloat16) if pixel_values is not None else None
+        ],dim=0)
+    if grid_thws is not None:
+        grid_thws = torch.cat([
+            grid_thws.to(device=visual_tokenizer.device) if grid_thws is not None else None
+        ],dim=0)
+    return input_ids, pixel_values, attention_mask, grid_thws, vae_pixel_values
+
+def pipe_img_edit(model, input_img, prompt, steps, txt_cfg, img_cfg, seed=42):
+    text_tokenizer = model.get_text_tokenizer()
+    visual_tokenizer = model.get_visual_tokenizer()
+    
+    width, height = input_img.size
+    height, width = visual_tokenizer.smart_resize(height, width, factor=32)
+
+    gen_kwargs = dict(
+          max_new_tokens=1024,
+          do_sample=False,
+          top_p=None,
+          top_k=None,
+          temperature=None,
+          repetition_penalty=None,
+          eos_token_id=text_tokenizer.eos_token_id,
+          pad_token_id=text_tokenizer.pad_token_id,
+          use_cache=True,
+          height=height,
+          width=width,
+          num_steps=steps,
+          seed=seed,
+          img_cfg=img_cfg,
+          txt_cfg=txt_cfg,
+      )
+    uncond_image = load_blank_image(width, height)
+    uncond_prompt = "<image>\nGenerate an image."
+    input_ids, pixel_values, attention_mask, grid_thws, _ = build_inputs(model, text_tokenizer, visual_tokenizer, uncond_prompt, uncond_image, width, height)
+    with torch.inference_mode():
+        no_both_cond = model.generate_condition(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs)
+
+    input_img = input_img.resize((width, height))
+    prompt = "<image>\n" + prompt.strip()
+    with torch.inference_mode():
+        input_ids, pixel_values, attention_mask, grid_thws, _ = build_inputs(model, text_tokenizer, visual_tokenizer, uncond_prompt, input_img, width, height)
+        no_txt_cond = model.generate_condition(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs) 
+
+    input_ids, pixel_values, attention_mask, grid_thws, vae_pixel_values = build_inputs(model, text_tokenizer, visual_tokenizer, prompt, input_img, width, height)
+    with torch.inference_mode():
+        cond = model.generate_condition(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs)
+        cond["vae_pixel_values"] = vae_pixel_values
+        images = model.generate_img(cond=cond, no_both_cond=no_both_cond, no_txt_cond=no_txt_cond, **gen_kwargs)
+    return images
+
+def main():
+    args = parse_args()
+    model, loading_info = AutoModelForCausalLM.from_pretrained(args.model_path,
+                                                torch_dtype=torch.bfloat16,
+                                                output_loading_info=True,
+                                                trust_remote_code=True
+                                                )
+    print(f'Loading info of Ovis-U1:\n{loading_info}')
+
+    model = model.eval().to("cuda")
+    model = model.to(torch.bfloat16)
+    image_path = os.path.join(os.path.dirname(__file__), "docs", "imgs", "cat.png")
+    pil_img = Image.open(image_path).convert('RGB')
+    prompt = "add a hat to this cat."
+    image = pipe_img_edit(model, pil_img, prompt, 
+                          args.steps, args.txt_cfg, args.img_cfg)[0]
+    image.save("test_image_edit.png")
+
+
+if __name__ == "__main__":
+    main()

test_img_to_txt.py

@@ -0,0 +1,84 @@
+import os
+import argparse
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Test Text Generation")
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="AIDC-AI/Ovis-U1-3B",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def build_inputs(model, text_tokenizer, visual_tokenizer, prompt, pil_image):
+    prompt, input_ids, pixel_values, grid_thws = model.preprocess_inputs(
+        prompt, 
+        [pil_image], 
+        generation_preface=None,
+        return_labels=False,
+        propagate_exception=False,
+        multimodal_type='single_image',
+        fix_sample_overall_length_navit=False
+        )
+    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
+    input_ids = input_ids.unsqueeze(0).to(device=model.device)
+    attention_mask = attention_mask.unsqueeze(0).to(device=model.device)
+    if pixel_values is not None:
+        pixel_values = torch.cat([
+            pixel_values.to(device=visual_tokenizer.device, dtype=torch.bfloat16) if pixel_values is not None else None
+        ],dim=0)
+    if grid_thws is not None:
+        grid_thws = torch.cat([
+            grid_thws.to(device=visual_tokenizer.device) if grid_thws is not None else None
+        ],dim=0)
+    return input_ids, pixel_values, attention_mask, grid_thws
+
+
+def pipe_txt_gen(model, pil_image, prompt):
+    text_tokenizer = model.get_text_tokenizer()
+    visual_tokenizer = model.get_visual_tokenizer()
+    gen_kwargs = dict(
+          max_new_tokens=4096,
+          do_sample=False,
+          top_p=None,
+          top_k=None,
+          temperature=None,
+          repetition_penalty=None,
+          eos_token_id=text_tokenizer.eos_token_id,
+          pad_token_id=text_tokenizer.pad_token_id,
+          use_cache=True,
+      )
+    prompt = "<image>\n" + prompt
+    input_ids, pixel_values, attention_mask, grid_thws = build_inputs(model, text_tokenizer, visual_tokenizer, prompt, pil_image)
+    with torch.inference_mode():
+        output_ids = model.generate(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs)[0]
+        gen_text = text_tokenizer.decode(output_ids, skip_special_tokens=True)
+    return gen_text
+
+
+def main():
+    # load model
+    args = parse_args()
+    model, loading_info = AutoModelForCausalLM.from_pretrained(args.model_path,
+                                                torch_dtype=torch.bfloat16,
+                                                output_loading_info=True,
+                                                trust_remote_code=True
+                                                )
+    print(f'Loading info of Ovis-U1:\n{loading_info}')
+
+    model = model.eval().to("cuda")
+    model = model.to(torch.bfloat16)
+    image_path = os.path.join(os.path.dirname(__file__), "docs", "imgs", "cat.png")
+    pil_img = Image.open(image_path).convert('RGB')
+    prompt = "What is it?"
+    gen_txt = pipe_txt_gen(model, pil_img, prompt)
+    print(gen_txt)
+
+
+if __name__ == "__main__":
+    main()

test_txt_to_img.py

@@ -0,0 +1,132 @@
+import os
+import argparse
+import math
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Test Text-to-Image")
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="AIDC-AI/Ovis-U1-3B",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=1024,
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=1024,
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42,
+    )
+    parser.add_argument(
+        "--steps", type=int, default=50,
+    )
+    parser.add_argument(
+        "--txt_cfg", type=float, default=5,
+    )
+    args = parser.parse_args()
+    return args
+
+
+def load_blank_image(width, height):
+    pil_image = Image.new("RGB", (width, height), (255, 255, 255)).convert('RGB')
+    return pil_image
+
+def build_inputs(model, text_tokenizer, visual_tokenizer, prompt, pil_image, target_width, target_height):
+    if pil_image is not None:
+        target_size = (int(target_width), int(target_height))
+        pil_image, vae_pixel_values, cond_img_ids = model.visual_generator.process_image_aspectratio(pil_image, target_size)
+        cond_img_ids[..., 0] = 1.0
+        vae_pixel_values = vae_pixel_values.unsqueeze(0).to(device=model.device)
+        width = pil_image.width
+        height = pil_image.height
+        resized_height, resized_width = visual_tokenizer.smart_resize(height, width, max_pixels=visual_tokenizer.image_processor.min_pixels)
+        pil_image = pil_image.resize((resized_width, resized_height))
+    else:
+        vae_pixel_values = None
+        cond_img_ids = None
+
+    prompt, input_ids, pixel_values, grid_thws = model.preprocess_inputs(
+        prompt, 
+        [pil_image], 
+        generation_preface=None,
+        return_labels=False,
+        propagate_exception=False,
+        multimodal_type='single_image',
+        fix_sample_overall_length_navit=False
+        )
+    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
+    input_ids = input_ids.unsqueeze(0).to(device=model.device)
+    attention_mask = attention_mask.unsqueeze(0).to(device=model.device)
+    if pixel_values is not None:
+        pixel_values = torch.cat([
+            pixel_values.to(device=visual_tokenizer.device, dtype=torch.bfloat16) if pixel_values is not None else None
+        ],dim=0)
+    if grid_thws is not None:
+        grid_thws = torch.cat([
+            grid_thws.to(device=visual_tokenizer.device) if grid_thws is not None else None
+        ],dim=0)
+    return input_ids, pixel_values, attention_mask, grid_thws, vae_pixel_values
+
+
+def pipe_t2i(model, prompt, height, width, steps, cfg, seed=42):
+    text_tokenizer = model.get_text_tokenizer()
+    visual_tokenizer = model.get_visual_tokenizer()
+    gen_kwargs = dict(
+          max_new_tokens=1024,
+          do_sample=False,
+          top_p=None,
+          top_k=None,
+          temperature=None,
+          repetition_penalty=None,
+          eos_token_id=text_tokenizer.eos_token_id,
+          pad_token_id=text_tokenizer.pad_token_id,
+          use_cache=True,
+          height=height,
+          width=width,
+          num_steps=steps,
+          seed=seed,
+          img_cfg=0,
+          txt_cfg=cfg,
+      )
+    uncond_image = load_blank_image(width, height)
+    uncond_prompt = "<image>\nGenerate an image."
+    input_ids, pixel_values, attention_mask, grid_thws, _ = build_inputs(model, text_tokenizer, visual_tokenizer, uncond_prompt, uncond_image, width, height)
+    with torch.inference_mode():
+        no_both_cond = model.generate_condition(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs)
+    prompt = "<image>\nDescribe the image by detailing the color, shape, size, texture, quantity, text, and spatial relationships of the objects:" + prompt
+    no_txt_cond = None
+    input_ids, pixel_values, attention_mask, grid_thws, vae_pixel_values = build_inputs(model, text_tokenizer, visual_tokenizer, prompt, uncond_image, width, height)
+    with torch.inference_mode():
+        cond = model.generate_condition(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, grid_thws=grid_thws, **gen_kwargs)
+        cond["vae_pixel_values"] = vae_pixel_values
+        images = model.generate_img(cond=cond, no_both_cond=no_both_cond, no_txt_cond=no_txt_cond, **gen_kwargs)
+    return images
+
+
+def main():
+    args = parse_args()
+    model, loading_info = AutoModelForCausalLM.from_pretrained(args.model_path,
+                                                torch_dtype=torch.bfloat16,
+                                                output_loading_info=True,
+                                                trust_remote_code=True
+                                                )
+    print(f'Loading info of Ovis-U1:\n{loading_info}')
+
+    model = model.eval().to("cuda")
+    model = model.to(torch.bfloat16)
+    prompt = "a cute cat"
+    image = pipe_t2i(model, prompt, args.height, args.width, args.steps, args.txt_cfg)[0]
+    image.save("test_t2i.png")
+
+
+if __name__ == "__main__":
+    main()