Update app.py

app.py

@@ -12,30 +12,14 @@ NEGATIVE_PROMPT = '''
 low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.
 '''
 
-
-def parse_args():
-    args = argparse.ArgumentParser(description='MindOmni')
-    args.add_argument('--device', type=str, default='cuda')
-    args.add_argument('--dtype', type=str, default='bf16')
-    args.add_argument('--model_path', type=str,
-                      default='EasonXiao-888/MindOmni')
-    args = args.parse_args()
-    return args
-
-
-def build_model(args):
-    device = args.device
-    MindOmni_model = MindOmni.from_pretrained(args.model_path)
-    if args.dtype == "bf16":
-        dtype = torch.bfloat16
-    MindOmni_model.to(device=device, dtype=dtype)
-    MindOmni_model.eval()
-    return MindOmni_model
+MindOmni_model = MindOmni.from_pretrained('EasonXiao-888/MindOmni')
+MindOmni_model.to(device='cuda', dtype=torch.bfloat16)
+MindOmni_model.eval()
 
 
 @spaces.GPU
 def understand_func(
-        MindOmni_model, text, do_sample, temperature,
+        text, do_sample, temperature,
         max_new_tokens, input_llm_images):
     if input_llm_images is not None and not isinstance(input_llm_images, list):
         input_llm_images = [input_llm_images]
@@ -47,7 +31,7 @@ def understand_func(
 
 @spaces.GPU
 def generate_func(
-        MindOmni_model, text, use_cot, height, width, guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model, max_input_image_size, randomize_seed, save_images, do_sample, temperature, max_new_tokens, input_llm_images, only_understand):
+        text, use_cot, height, width, guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model, max_input_image_size, randomize_seed, save_images, do_sample, temperature, max_new_tokens, input_llm_images, only_understand):
     if input_llm_images is not None and not isinstance(input_llm_images, list):
         input_llm_images = [input_llm_images]
 
@@ -76,7 +60,7 @@ def generate_func(
     return img, prompt_, seed
 
 
-def build_gradio(args, MindOmni_model):
+def build_gradio():
     with gr.Blocks() as demo:
         gr.Markdown("## 🪄 MindOmni Demo")
 
@@ -133,7 +117,7 @@ def build_gradio(args, MindOmni_model):
                             )
 
                 g_btn.click(
-                    partial(generate_func, MindOmni_model),
+                    generate_func,
                     inputs=[g_prompt, g_use_cot, g_height, g_width, g_scale, g_steps,
                             g_seed, g_sep_cfg, g_offload, g_max_img, g_rand, g_save,
                             g_do_sample, g_temperature, g_max_new_tok,
@@ -156,7 +140,7 @@ def build_gradio(args, MindOmni_model):
                         u_answer = gr.Textbox(label="Answer", lines=8)
 
                 u_btn.click(
-                    partial(understand_func, MindOmni_model),
+                    understand_func,
                     inputs=[u_prompt, u_do_sample,
                             u_temperature, u_max_new_tok, u_image],
                     outputs=u_answer)
@@ -164,12 +148,5 @@ def build_gradio(args, MindOmni_model):
         demo.launch()
 
 
-def main():
-    args = parse_args()
-    print(f'running args: {args}')
-    MindOmni_model = build_model(args)
-    build_gradio(args, MindOmni_model)
-
-
 if __name__ == '__main__':
-    main()
+    build_gradio()