Update app.py

app.py

@@ -31,7 +31,7 @@ def understand_func(
 
 @spaces.GPU
 def generate_func(
-        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):
+        MindOmni_model, text, use_cot, cascade_thinking, 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]
 
@@ -41,7 +41,8 @@ def generate_func(
     print(f'Generate image prompt: {text}')
     output, prompt_ = MindOmni_model.generate_image(
         height, width, guidance_scale, inference_steps, separate_cfg_infer, offload_model, seed, max_input_image_size,
-        text, NEGATIVE_PROMPT, input_llm_images, do_sample, temperature, max_new_tokens, only_understand, use_cot=use_cot)
+        text, NEGATIVE_PROMPT, input_llm_images, do_sample, temperature, max_new_tokens, only_understand, use_cot=use_cot,
+        cascade_thinking=cascade_thinking)
     print('Generation finished.')
     
     img = output[0]
@@ -74,8 +75,9 @@ def build_gradio():
                         g_btn = gr.Button("🚀 Generate Image")
 
                         with gr.Accordion("📚 Image Generation Args"):
-                            g_use_cot = gr.Checkbox(label="With thinking", value=False)
-                            g_do_sample = gr.Checkbox(label="Do sample", value=False)
+                            g_use_cot = gr.Checkbox(label="Use thinking", value=False)
+                            g_cascade_thinking = gr.Checkbox(label="Cascade thinking (experimental for better quality)", value=False)
+                            g_do_sample = gr.Checkbox(label="Do sample (for more diversity)", value=False)
                             g_temperature = gr.Slider(0, 10, value=1, label="Temperature")
                             g_max_new_tok = gr.Slider(32, 8192, value=512, label="Max new tokens")
 
@@ -99,25 +101,25 @@ def build_gradio():
                         with gr.Accordion("🖼️ Prompt Examples: Text-only"):
                             gr.Examples(
                                 examples=[
-                                    ["Futuristic city skyline at sunset, digital art", 42, False, False, False, 1024, 1024, "assets/example_outputs/case_1.png"],
-                                    ["An image of China's national treasure animal.", 42, False, True, False, 1024, 1024, "assets/example_outputs/case_2.png"],
-                                    ["Scene in the Sydney Opera House when New York is at noon.", 42, False, True, False, 1024, 1024, "assets/example_outputs/case_3.png"],
-                                    ["Generate an image of an animal with (3 + 6) lives", 7393438, False, True, False, 1024, 1024, "assets/example_outputs/case_4.png"],
+                                    ["Futuristic city skyline at sunset, digital art", 42, False, False, False, False, 1024, 1024, "assets/example_outputs/case_1.png"],
+                                    ["An image of China's national treasure animal.", 42, False, True, False, False, 1024, 1024, "assets/example_outputs/case_2.png"],
+                                    ["Scene in the Sydney Opera House when New York is at noon.", 42, False, True, False, False, 1024, 1024, "assets/example_outputs/case_3.png"],
+                                    ["Generate an image of an animal with (3 + 6) lives", 7393438, False, True, False, False, 1024, 1024, "assets/example_outputs/case_4.png"],
                                 ],
-                                inputs=[g_prompt, g_seed, g_rand, g_use_cot, g_do_sample, g_height, g_width, g_out_img],
+                                inputs=[g_prompt, g_seed, g_rand, g_use_cot, g_cascade_thinking, g_do_sample, g_height, g_width, g_out_img],
                             )
                         with gr.Accordion("🖼️ Prompt Examples: With reference image"):
                             gr.Examples(
                                 examples=[
-                                    ["An image of the animal growing up", "assets/tapdole.jpeg", 42, False, True, True, 1024, 1024, "assets/example_outputs/case_5.png"]
+                                    ["An image of the animal growing up", "assets/tapdole.jpeg", 42, False, True, False, True, 1024, 1024, "assets/example_outputs/case_5.png"]
                                 ],
-                                inputs=[g_prompt, g_image, g_seed, g_rand, g_use_cot, g_do_sample, g_height, g_width, g_out_img],
+                                inputs=[g_prompt, g_image, g_seed, g_rand, g_use_cot, g_cascade_thinking, g_do_sample, g_height, g_width, g_out_img],
                             )
 
                 g_btn.click(
                     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,
+                    inputs=[g_prompt, g_use_cot, g_cascade_thinking, 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,
                             g_image, gr.State(False)],          # only_understand=False
                     outputs=[g_out_img, g_prompt_out, g_seed_out])