evalstate/Ovis-U1-3B
0
1import os2import subprocess3subprocess.run('pip install flash-attn==2.6.3 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)4import random5import spaces6import numpy as np7import torch8from PIL import Image9import gradio as gr10from transformers import AutoModelForCausalLM11from test_img_edit import pipe_img_edit12from test_img_to_txt import pipe_txt_gen13from test_txt_to_img import pipe_t2i14 15 16# Constants17MAX_SEED = 1000018 19hf_token = os.getenv("HF_TOKEN")20 21HUB_MODEL_ID = "AIDC-AI/Ovis-U1-3B"22model, loading_info = AutoModelForCausalLM.from_pretrained(23 HUB_MODEL_ID, 24 torch_dtype=torch.bfloat16,25 output_loading_info=True,26 token=hf_token,27 trust_remote_code=True28 )29print(f'Loading info of Ovis-U1:\n{loading_info}')30 31model = model.eval().to("cuda")32model = model.to(torch.bfloat16)33 34def set_global_seed(seed: int = 42,show_api=False):35 random.seed(seed)36 np.random.seed(seed)37 torch.manual_seed(seed)38 torch.cuda.manual_seed_all(seed)39 40def randomize_seed_fn(seed: int, randomize: bool,show_api=False) -> int:41 return random.randint(0, MAX_SEED) if randomize else seed42 43@spaces.GPU(duration=20)44def 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]:45 """Use Ovis-U1-3B to generate an image. Supply a Text Prompt"""46 set_global_seed(final_seed)47 images = pipe_t2i(model, prompt, height, width, steps, cfg=guidance_scale, seed=final_seed)48 return images49 50@spaces.GPU(duration=20)51def process_img_to_txt(prompt: str, img: Image.Image, progress: gr.Progress = gr.Progress(track_tqdm=True)) -> str:52 """Use Ovis-U1-3B to analyse an Image"""53 output_text = pipe_txt_gen(model, img, prompt)54 return output_text55 56@spaces.GPU(duration=20)57def 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]:58 """Use Ovis-U1-3B to modify an Image. Supply an Image URL and a Text Prompt (e.g. 'ghiblify', 'low-poly 3d render', 'replace house with car'"""59 set_global_seed(final_seed)60 images = pipe_img_edit(model, img, prompt, steps, txt_cfg, img_cfg, seed=final_seed)61 return images62 63# Gradio UI64with gr.Blocks(title="Ovis-U1-3B") as demo:65 gr.Markdown('''# Ovis-U1-3B66 ''')67 68 with gr.Row():69 with gr.Column():70 with gr.Tabs():71 with gr.TabItem("Image + Text → Image"):72 edit_image_input = gr.Image(label="Input Image", type="pil")73 with gr.Row():74 edit_prompt_input = gr.Textbox(75 label="Prompt",76 show_label=False,77 placeholder="Describe the editing instruction...",78 container=False,79 lines=180 )81 run_edit_image_btn = gr.Button("Run", scale=0)82 83 with gr.Accordion("Advanced Settings", open=False):84 85 with gr.Row():86 87 edit_img_guidance_slider = gr.Slider(88 label="Image Guidance Scale",89 minimum=1.0, maximum=10.0,90 step=0.1, value=1.591 )92 93 edit_txt_guidance_slider = gr.Slider(94 label="Text Guidance Scale",95 minimum=1.0, maximum=30.0,96 step=0.5, value=6.097 )98 99 edit_num_steps_slider = gr.Slider(100 label='Steps',101 minimum=40, maximum=100, 102 value=50, step=1103 )104 edit_seed_slider = gr.Slider(105 label="Seed",106 minimum=0, maximum=int(MAX_SEED),107 step=1, value=42108 )109 edit_randomize_checkbox = gr.Checkbox(110 label="Randomize seed", value=False111 )112 113 img_edit_examples_data = [114 ["imgs/train.png", "Modify this image in a Ghibli style. "],115 ["imgs/chair.png", "Transfer the image into a faceted low-poly 3-D render style."],116 ["imgs/car.png", "Replace the tiny house on wheels in the image with a vintage car."],117 ]118 gr.Examples(119 examples=img_edit_examples_data,120 inputs=[edit_image_input, edit_prompt_input],121 cache_examples=False, 122 label="Image Editing Examples"123 )124 125 with gr.TabItem("Text → Image"):126 with gr.Row():127 prompt_gen_input = gr.Textbox(128 label="Prompt",129 show_label=False,130 placeholder="Describe the image you want...",131 container=False,132 lines=1133 )134 run_image_gen_btn = gr.Button("Run", scale=0)135 136 with gr.Accordion("Advanced Settings", open=False):137 with gr.Row():138 height_slider = gr.Slider(139 label='height', 140 minimum=256, maximum=1536, 141 value=1024, step=32142 )143 width_slider = gr.Slider(144 label='width', 145 minimum=256, maximum=1536, 146 value=1024, step=32147 )148 149 guidance_slider = gr.Slider(150 label="Guidance Scale",151 minimum=1.0, maximum=30.0,152 step=0.5, value=5.0153 )154 155 num_steps_slider = gr.Slider(156 label='Steps',157 minimum=40, maximum=100, 158 value=50, step=1159 )160 seed_slider = gr.Slider(161 label="Seed",162 minimum=0, maximum=int(MAX_SEED),163 step=1, value=42164 )165 randomize_checkbox = gr.Checkbox(166 label="Randomize seed", value=False167 )168 169 text_gen_examples_data = [170 ["A breathtaking fairy with teal wings sits gracefully on a lotus flower in a serene pond, exuding elegance."],171 ["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."],172 ["A photo of a pug wearing a cowboy hat and bandana, sitting on a hay bale."]173 ]174 gr.Examples(175 examples=text_gen_examples_data,176 inputs=[prompt_gen_input],177 cache_examples=False, 178 label="Image Generation Examples"179 )180 181 with gr.TabItem("Image → Text"):182 image_understand_input = gr.Image(label="Input Image", type="pil")183 with gr.Row():184 prompt_understand_input = gr.Textbox(185 label="Prompt",186 show_label=False,187 placeholder="Describe the question about image...",188 container=False,189 lines=1190 )191 run_image_understand_btn = gr.Button("Run", scale=0)192 193 image_understanding_examples_data = [194 ["imgs/table.webp", "In what scenario does this picture take place?"],195 ["imgs/count.png", "How many broccoli are there in the picture?"],196 ["imgs/foot.webp", "Where is this picture located?"],197 ]198 gr.Examples(199 examples=image_understanding_examples_data,200 inputs=[image_understand_input, prompt_understand_input],201 cache_examples=False, 202 label="Image Understanding Examples"203 )204 205 clean_btn = gr.Button("Clear All Inputs/Outputs")206 207 with gr.Column():208 output_gallery = gr.Gallery(label="Generated Images", columns=2, visible=True) # Default to visible, content will control209 output_text = gr.Textbox(label="Generated Text", visible=False, lines=5, interactive=False)210 211 @spaces.GPU(duration=20)212 def run_img_txt_to_img_tab(prompt, img, steps, seed, txt_cfg, img_cfg, progress=gr.Progress(track_tqdm=True),show_api=False):213 if img is None:214 return (215 gr.update(value=[], visible=False),216 gr.update(value="Please upload an image for editing.", visible=True)217 )218 # Seed is already finalized by the randomize_seed_fn in the click chain219 imgs = process_img_txt_to_img(prompt, img, steps, seed, txt_cfg, img_cfg, progress=progress)220 return (221 gr.update(value=imgs, visible=True),222 gr.update(value="", visible=False)223 )224 225 @spaces.GPU(duration=20)226 def run_txt_to_img_tab(prompt, height, width, steps, seed, guidance, progress=gr.Progress(track_tqdm=True),show_api=False):227 """Use Ovis-U1-3B to generate an Image."""228 # Seed is already finalized by the randomize_seed_fn in the click chain229 imgs = process_txt_to_img(prompt, height, width, steps, seed, guidance, progress=progress)230 return (231 gr.update(value=imgs, visible=True),232 gr.update(value="", visible=False)233 )234 235 @spaces.GPU(duration=20)236 def run_img_to_txt_tab(img, prompt, progress=gr.Progress(track_tqdm=True),show_api=False):237 """Use Ovis-U1-3B to understand an Image (Vision processing)"""238 if img is None:239 return (240 gr.update(value=[], visible=False),241 gr.update(value="Please upload an image for understanding.", visible=True)242 )243 txt = process_img_to_txt(prompt, img, progress=progress)244 return (245 gr.update(value=[], visible=False),246 gr.update(value=txt, visible=True)247 )248 249 def clean_all_fn():250 return (251 # Tab 1 inputs252 gr.update(value=None),253 gr.update(value=""),254 gr.update(value=1.5),255 gr.update(value=6.0),256 gr.update(value=50),257 gr.update(value=42),258 gr.update(value=False),259 # Tab 2 inputs260 gr.update(value=""), # prompt_gen_input261 gr.update(value=1024),262 gr.update(value=1024),263 gr.update(value=5.0),264 gr.update(value=50),265 gr.update(value=42), # seed_slider266 gr.update(value=False), # randomize_checkbox267 # Tab 3 inputs268 gr.update(value=None), # image_understand_input269 gr.update(value=""), # prompt_understand_input270 # Outputs271 gr.update(value=[], visible=True), # output_gallery (reset and keep visible for next gen)272 gr.update(value="", visible=False) # output_text (reset and hide)273 )274 275 # Event listeners for Image + Text -> Image276 edit_inputs = [edit_prompt_input, edit_image_input, edit_num_steps_slider, edit_seed_slider, edit_txt_guidance_slider, edit_img_guidance_slider]277 278 run_edit_image_btn.click(279 fn=randomize_seed_fn,280 inputs=[edit_seed_slider, edit_randomize_checkbox],281 outputs=[edit_seed_slider]282 ).then(283 fn=run_img_txt_to_img_tab,284 inputs=edit_inputs, 285 outputs=[output_gallery, output_text]286 )287 288 edit_prompt_input.submit(289 fn=randomize_seed_fn,290 inputs=[edit_seed_slider, edit_randomize_checkbox],291 outputs=[edit_seed_slider]292 ).then(293 fn=run_img_txt_to_img_tab,294 inputs=edit_inputs,295 outputs=[output_gallery, output_text]296 )297 298 # Event listeners for Text -> Image299 # gen_inputs = [prompt_gen_input, height_slider, width_slider, num_steps_slider, seed_slider, guidance_slider]300 301 # run_image_gen_btn.click(302 # fn=randomize_seed_fn,303 # inputs=[seed_slider, randomize_checkbox],304 # outputs=[seed_slider]305 # ).then(306 # fn=run_txt_to_img_tab,307 # inputs=gen_inputs, 308 # outputs=[output_gallery, output_text]309 # )310 311 # prompt_gen_input.submit(312 # fn=randomize_seed_fn,313 # inputs=[seed_slider, randomize_checkbox],314 # outputs=[seed_slider]315 # ).then(316 # fn=run_txt_to_img_tab,317 # inputs=gen_inputs,318 # outputs=[output_gallery, output_text]319 # )320 321 # Event listeners for Image -> Text322 # understand_inputs = [image_understand_input, prompt_understand_input]323 324 # run_image_understand_btn.click(325 # fn=run_img_to_txt_tab,326 # inputs=understand_inputs,327 # outputs=[output_gallery, output_text]328 # )329 330 # prompt_understand_input.submit(331 # fn=run_img_to_txt_tab,332 # inputs=understand_inputs,333 # outputs=[output_gallery, output_text]334 # )335 336 # clean_btn.click(337 # fn=clean_all_fn,338 # inputs=[],339 # outputs=[340 # edit_image_input, edit_prompt_input, edit_img_guidance_slider, edit_txt_guidance_slider,341 # edit_num_steps_slider, edit_seed_slider, edit_randomize_checkbox,342 # prompt_gen_input, height_slider, width_slider, guidance_slider, num_steps_slider, seed_slider, randomize_checkbox, 343 # image_understand_input, prompt_understand_input,344 # output_gallery, output_text345 # ]346 # )347 348if __name__ == "__main__":349 demo.launch(mcp_server=True)