Update app.py

app.py

@@ -13,7 +13,6 @@ subprocess.run(
 
 from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights
 from PIL import Image
-import uuid
 
 from data.data_utils import add_special_tokens, pil_img2rgb
 from data.transforms import ImageTransform
@@ -32,19 +31,16 @@ save_dir = "./model"
 repo_id = "ByteDance-Seed/BAGEL-7B-MoT"
 cache_dir = save_dir + "/cache"
 
-if not os.path.exists(os.path.join(save_dir, "ema.safetensors")):
-    print(f"Downloading model from {repo_id} to {save_dir}")
-    snapshot_download(cache_dir=cache_dir,
-                      local_dir=save_dir,
-                      repo_id=repo_id,
-                      local_dir_use_symlinks=False,
-                      resume_download=True,
-                      allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
-                      )
-else:
-    print(f"Model found at {save_dir}")
+snapshot_download(cache_dir=cache_dir,
+  local_dir=save_dir,
+  repo_id=repo_id,
+  local_dir_use_symlinks=False,
+  resume_download=True,
+  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
+)
 
-model_path = "./model"
+# Model Initialization
+model_path = "./model" #Download from https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT
 
 llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
 llm_config.qk_norm = True
@@ -60,7 +56,7 @@ vae_model, vae_config = load_ae(local_path=os.path.join(model_path, "ae.safetens
 config = BagelConfig(
     visual_gen=True,
     visual_und=True,
-    llm_config=llm_config,
+    llm_config=llm_config, 
     vit_config=vit_config,
     vae_config=vae_config,
     vit_max_num_patch_per_side=70,
@@ -81,6 +77,7 @@ tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
 vae_transform = ImageTransform(1024, 512, 16)
 vit_transform = ImageTransform(980, 224, 14)
 
+# Model Loading and Multi GPU Infernece Preparing
 device_map = infer_auto_device_map(
     model,
     max_memory={i: "80GiB" for i in range(torch.cuda.device_count())},
@@ -100,20 +97,16 @@ same_device_modules = [
 if torch.cuda.device_count() == 1:
     first_device = device_map.get(same_device_modules[0], "cuda:0")
     for k in same_device_modules:
-        device_map[k] = first_device
+        if k in device_map:
+            device_map[k] = first_device
+        else:
+            device_map[k] = "cuda:0"
 else:
-    # Ensure all same_device_modules are on the same device if they exist in device_map
-    # Find the device for the first module in the list that is actually in the device_map
-    first_assigned_device = None
-    for k_module in same_device_modules:
-        if k_module in device_map:
-            first_assigned_device = device_map[k_module]
-            break
-    if first_assigned_device is not None:
-        for k_module in same_device_modules:
-            if k_module in device_map: # Only assign if the module is part of the device_map
-                device_map[k_module] = first_assigned_device
-
+    first_device = device_map.get(same_device_modules[0])
+    for k in same_device_modules:
+        if k in device_map:
+            device_map[k] = first_device
+            
 model = load_checkpoint_and_dispatch(
     model,
     checkpoint=os.path.join(model_path, "ema.safetensors"),
@@ -123,6 +116,8 @@ model = load_checkpoint_and_dispatch(
     force_hooks=True,
 ).eval()
 
+
+# Inferencer Preparing 
 inferencer = InterleaveInferencer(
     model=model,
     vae_model=vae_model,
@@ -133,7 +128,8 @@ inferencer = InterleaveInferencer(
 )
 
 def set_seed(seed):
-    if seed is not None and seed > 0:
+    """Set random seeds for reproducibility"""
+    if seed > 0:
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
@@ -144,384 +140,389 @@ def set_seed(seed):
         torch.backends.cudnn.benchmark = False
     return seed
 
-# --- Backend Functions (Adapted from original app.py) ---
-@spaces.GPU(duration=90)
-def call_text_to_image(prompt, show_thinking, cfg_text_scale, cfg_interval, 
-                       timestep_shift, num_timesteps, cfg_renorm_min, cfg_renorm_type, 
-                       max_think_token_n, do_sample, text_temperature, seed, image_ratio):
+# Text to Image function with thinking option and hyperparameters
+def text_to_image(prompt, show_thinking=False, cfg_text_scale=4.0, cfg_interval=0.4, 
+                 timestep_shift=3.0, num_timesteps=50, 
+                 cfg_renorm_min=1.0, cfg_renorm_type="global", 
+                 max_think_token_n=1024, do_sample=False, text_temperature=0.3,
+                 seed=0, image_ratio="1:1"):
+    # Set seed for reproducibility
     set_seed(seed)
-    image_shapes = (1024, 1024)
-    if image_ratio == "4:3": image_shapes = (768, 1024)
-    elif image_ratio == "3:4": image_shapes = (1024, 768)
-    elif image_ratio == "16:9": image_shapes = (576, 1024)
-    elif image_ratio == "9:16": image_shapes = (1024, 576)
 
+    if image_ratio == "1:1":
+        image_shapes = (1024, 1024)
+    elif image_ratio == "4:3":
+        image_shapes = (768, 1024)
+    elif image_ratio == "3:4":
+        image_shapes = (1024, 768) 
+    elif image_ratio == "16:9":
+        image_shapes = (576, 1024)
+    elif image_ratio == "9:16":
+        image_shapes = (1024, 576) 
+    
+    # Set hyperparameters
     inference_hyper = dict(
         max_think_token_n=max_think_token_n if show_thinking else 1024,
         do_sample=do_sample if show_thinking else False,
         text_temperature=text_temperature if show_thinking else 0.3,
         cfg_text_scale=cfg_text_scale,
-        cfg_interval=[cfg_interval, 1.0],
+        cfg_interval=[cfg_interval, 1.0],  # End fixed at 1.0
         timestep_shift=timestep_shift,
         num_timesteps=num_timesteps,
         cfg_renorm_min=cfg_renorm_min,
         cfg_renorm_type=cfg_renorm_type,
         image_shapes=image_shapes,
     )
+    
+    # Call inferencer with or without think parameter based on user choice
     result = inferencer(text=prompt, think=show_thinking, **inference_hyper)
-    return result.get("image", None), result.get("text", None) # text is thinking
+    return result["image"], result.get("text", None)
+
+
+# Image Understanding function with thinking option and hyperparameters
+def image_understanding(image: Image.Image, prompt: str, show_thinking=False, 
+                        do_sample=False, text_temperature=0.3, max_new_tokens=512):
+    if image is None:
+        return "Please upload an image."
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
-@spaces.GPU(duration=90)
-def call_image_understanding(image, prompt, show_thinking, do_sample, text_temperature, max_new_tokens, seed):
-    set_seed(seed)
-    if image is None: return "Please upload an image.", None
-    if isinstance(image, np.ndarray): image = Image.fromarray(image)
     image = pil_img2rgb(image)
     
+    # Set hyperparameters
     inference_hyper = dict(
         do_sample=do_sample,
         text_temperature=text_temperature,
-        max_think_token_n=max_new_tokens, 
+        max_think_token_n=max_new_tokens, # Set max_length
     )
-    result = inferencer(image=image, text=prompt, think=show_thinking, understanding_output=True, **inference_hyper)
-    return result.get("text", None), None # Main output is text, thinking is part of it if show_thinking=True
-
-@spaces.GPU(duration=90)
-def call_edit_image(image, prompt, show_thinking, cfg_text_scale, cfg_img_scale, cfg_interval, 
-                    timestep_shift, num_timesteps, cfg_renorm_min, cfg_renorm_type, 
-                    max_think_token_n, do_sample, text_temperature, seed):
+    
+    # Use show_thinking parameter to control thinking process
+    result = inferencer(image=image, text=prompt, think=show_thinking, 
+                        understanding_output=True, **inference_hyper)
+    return result["text"]
+
+
+# Image Editing function with thinking option and hyperparameters
+def edit_image(image: Image.Image, prompt: str, show_thinking=False, cfg_text_scale=4.0, 
+              cfg_img_scale=2.0, cfg_interval=0.0, 
+              timestep_shift=3.0, num_timesteps=50, cfg_renorm_min=1.0, 
+              cfg_renorm_type="text_channel", max_think_token_n=1024, 
+              do_sample=False, text_temperature=0.3, seed=0):
+    # Set seed for reproducibility
     set_seed(seed)
-    if image is None: return "Please upload an image.", None, None
-    if isinstance(image, np.ndarray): image = Image.fromarray(image)
-    image = pil_img2rgb(image)
+    
+    if image is None:
+        return "Please upload an image.", ""
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
+    image = pil_img2rgb(image)
+    
+    # Set hyperparameters
     inference_hyper = dict(
         max_think_token_n=max_think_token_n if show_thinking else 1024,
         do_sample=do_sample if show_thinking else False,
         text_temperature=text_temperature if show_thinking else 0.3,
         cfg_text_scale=cfg_text_scale,
         cfg_img_scale=cfg_img_scale,
-        cfg_interval=[cfg_interval, 1.0],
+        cfg_interval=[cfg_interval, 1.0],  # End fixed at 1.0
         timestep_shift=timestep_shift,
         num_timesteps=num_timesteps,
         cfg_renorm_min=cfg_renorm_min,
         cfg_renorm_type=cfg_renorm_type,
     )
-    result = inferencer(image=image, text=prompt, think=show_thinking, **inference_hyper)
-    return result.get("image", None), result.get("text", None) # text is thinking
-
-# --- Gradio UI --- 
-
-DEFAULT_WELCOME_MESSAGE = {
-    "role": "assistant",
-    "content": "Hello! I am BAGEL, your multimodal assistant. How can I help you today? Select a mode and enter your prompt.",
-    "key": "welcome"
-}
-
-class GradioApp:
-    def __init__(self):
-        self.current_conversation_id = None
-        self.conversation_contexts = {}
-        self.conversations_list = [] # For the sidebar
-
-    def _get_current_history(self):
-        if self.current_conversation_id and self.current_conversation_id in self.conversation_contexts:
-            return self.conversation_contexts[self.current_conversation_id]["history"]
-        return []
-
-    def _get_current_settings(self):
-        if self.current_conversation_id and self.current_conversation_id in self.conversation_contexts:
-            return self.conversation_contexts[self.current_conversation_id].get("settings", {})
-        return {}
     
-    def _update_conversation_list_ui(self):
-        return gr.update(choices=[(c['label'], c['key']) for c in self.conversations_list], value=self.current_conversation_id)
-
-    def add_message(self, text_input, image_input, mode, 
-                    # TTI params
-                    tti_show_thinking, tti_cfg_text_scale, tti_cfg_interval, tti_timestep_shift, tti_num_timesteps, tti_cfg_renorm_min, tti_cfg_renorm_type, tti_max_think_token_n, tti_do_sample, tti_text_temperature, tti_seed, tti_image_ratio,
-                    # Edit params
-                    edit_show_thinking, edit_cfg_text_scale, edit_cfg_img_scale, edit_cfg_interval, edit_timestep_shift, edit_num_timesteps, edit_cfg_renorm_min, edit_cfg_renorm_type, edit_max_think_token_n, edit_do_sample, edit_text_temperature, edit_seed,
-                    # Understand params
-                    und_show_thinking, und_do_sample, und_text_temperature, und_max_new_tokens, und_seed
-                    ):
-        if not text_input and not (mode in ["Image Edit", "Image Understanding"] and image_input):
-            gr.Warning("Please enter a prompt or upload an image for Edit/Understanding modes.")
-            # Need to yield original state for all outputs if we return early
-            # This part is tricky with dynamic outputs, might need a dummy update for all
-            # For simplicity, let's assume user always provides some input
-            # A better way is to disable submit button if input is invalid
-            return self._get_current_history(), gr.update(value=None), gr.update(value=None) # chatbot, text_input, image_input
-
-        if not self.current_conversation_id:
-            self.new_chat_session(text_input[:30] if text_input else "New Chat") # Create a new chat if none exists
-
-        history = self._get_current_history()
-        
-        # Store settings for this turn
-        # This is simplified; best-gradio-ui.py stores settings per conversation
-        current_turn_settings = {
-            "mode": mode,
-            # Store PIL image directly if needed, or handle path carefully
-            "image_input": image_input, # Now storing the PIL image or None
-            # TTI
-            "tti_show_thinking": tti_show_thinking, "tti_cfg_text_scale": tti_cfg_text_scale, "tti_cfg_interval": tti_cfg_interval, "tti_timestep_shift": tti_timestep_shift, "tti_num_timesteps": tti_num_timesteps, "tti_cfg_renorm_min": tti_cfg_renorm_min, "tti_cfg_renorm_type": tti_cfg_renorm_type, "tti_max_think_token_n": tti_max_think_token_n, "tti_do_sample": tti_do_sample, "tti_text_temperature": tti_text_temperature, "tti_seed": tti_seed, "tti_image_ratio": tti_image_ratio,
-            # Edit
-            "edit_show_thinking": edit_show_thinking, "edit_cfg_text_scale": edit_cfg_text_scale, "edit_cfg_img_scale": edit_cfg_img_scale, "edit_cfg_interval": edit_cfg_interval, "edit_timestep_shift": edit_timestep_shift, "edit_num_timesteps": edit_num_timesteps, "edit_cfg_renorm_min": edit_cfg_renorm_min, "edit_cfg_renorm_type": edit_cfg_renorm_type, "edit_max_think_token_n": edit_max_think_token_n, "edit_do_sample": edit_do_sample, "edit_text_temperature": edit_text_temperature, "edit_seed": edit_seed,
-            # Understand
-            "und_show_thinking": und_show_thinking, "und_do_sample": und_do_sample, "und_text_temperature": und_text_temperature, "und_max_new_tokens": und_max_new_tokens, "und_seed": und_seed
-        }
-        self.conversation_contexts[self.current_conversation_id]["settings"] = current_turn_settings
-
-        user_content_list = []
-        if text_input:
-            user_content_list.append({"type": "text", "text": text_input})
-        if image_input and mode in ["Image Edit", "Image Understanding"]:
-            # For 'messages' format, images are typically handled by passing them as part of a list of content dicts.
-            # Gradio's Chatbot with type='messages' can render PIL Images or file paths directly in the 'content' list.
-            user_content_list.append({"type": "image", "image": image_input}) # Assuming image_input is PIL
-        
-        # Construct the user message for history
-        # If only text, content can be a string. If mixed, it's a list of dicts.
-        user_message_for_history = {
-            "role": "user", 
-            "content": text_input if not image_input else user_content_list, 
-            "key": str(uuid.uuid4())
-        }
-        if not text_input and image_input:
-             user_message_for_history["content"] = user_content_list
-        elif not user_content_list:
-            # Handle case where there's no input at all, though prior checks should prevent this.
-            gr.Warning("No input provided.")
-            return self._get_current_history(), gr.update(value=None), gr.update(value=None)
-
-
-        history.append(user_message_for_history)
-        history.append({"role": "assistant", "content": "Processing...", "key": str(uuid.uuid4())})
-        
-        yield history, gr.update(value=None), gr.update(value=None) # chatbot, text_input, image_input (clear inputs)
+    # Include thinking parameter based on user choice
+    result = inferencer(image=image, text=prompt, think=show_thinking, **inference_hyper)
+    return result["image"], result.get("text", "")
 
-        # Call backend
-        try:
-            output_image = None
-            output_text = None
-            thinking_text = None
-            
-            # image_input is already a PIL image from the gr.Image component with type="pil"
-            pil_image_input = image_input 
-
-            if mode == "Text to Image":
-                output_image, thinking_text = call_text_to_image(text_input, tti_show_thinking, tti_cfg_text_scale, tti_cfg_interval, tti_timestep_shift, tti_num_timesteps, tti_cfg_renorm_min, tti_cfg_renorm_type, tti_max_think_token_n, tti_do_sample, tti_text_temperature, tti_seed, tti_image_ratio)
-            elif mode == "Image Edit":
-                if not pil_image_input:
-                    output_text = "Error: Image required for Image Edit mode."
-                else:
-                    output_image, thinking_text = call_edit_image(pil_image_input, text_input, edit_show_thinking, edit_cfg_text_scale, edit_cfg_img_scale, edit_cfg_interval, edit_timestep_shift, edit_num_timesteps, edit_cfg_renorm_min, edit_cfg_renorm_type, edit_max_think_token_n, edit_do_sample, edit_text_temperature, edit_seed)
-            elif mode == "Image Understanding":
-                if not pil_image_input:
-                    output_text = "Error: Image required for Image Understanding mode."
-                else:
-                    output_text, _ = call_image_understanding(pil_image_input, text_input, und_show_thinking, und_do_sample, und_text_temperature, und_max_new_tokens, und_seed)
-                    # For VLM, the main output is text, thinking might be part of it or not separately returned
-                    # depending on `inferencer`'s behavior with `understanding_output=True`
-                    if und_show_thinking and output_text and "Thinking:" in output_text: # crude check
-                        parts = output_text.split("Thinking:", 1)
-                        if len(parts) > 1:
-                           thinking_text = "Thinking:" + parts[1].split("\nAnswer:")[0] if "\nAnswer:" in parts[1] else parts[1]
-                           output_text = parts[0].strip() + ("\nAnswer:" + output_text.split("\nAnswer:")[1] if "\nAnswer:" in output_text else "")
-                        else:
-                            thinking_text = None # Or handle as part of main output_text
-
-            bot_response_content = []
-            if thinking_text:
-                # For 'messages' type, each part of the content is a dict in a list
-                bot_response_content.append({"type": "text", "text": f"**Thinking Process:**\n{thinking_text}"}) 
-            if output_text:
-                bot_response_content.append({"type": "text", "text": output_text})
-            if output_image: # output_image should be a PIL Image
-                bot_response_content.append({"type": "image", "image": output_image})
-            
-            if not bot_response_content:
-                bot_response_content.append({"type": "text", "text": "(No output generated)"})
-
-            # Update the last message (which was "Processing...")
-            history[-1]["content"] = bot_response_content_list[0]["text"] if len(bot_response_content_list) == 1 and bot_response_content_list[0]["type"] == "text" else bot_response_content_list
-
-        except Exception as e:
-            print(f"Error during processing: {e}")
-            history[-1]["content"] = [{"type": "text", "content": f"Error: {str(e)}"}]
-            history[-1]["loading"] = False
-            raise gr.Error(f"Processing Error: {str(e)}")
-
-        yield history, gr.update(value=None), gr.update(value=None)
-
-    def new_chat_session(self, label="New Chat"):
-        session_id = str(uuid.uuid4())
-        self.current_conversation_id = session_id
-        self.conversation_contexts[session_id] = {
-            "history": [DEFAULT_WELCOME_MESSAGE.copy()],
-            "settings": {} # Initialize with default settings if any
-        }
-        # Ensure label is unique if needed, or just use the provided one
-        # For simplicity, we allow duplicate labels for now.
-        new_conv_entry = {"label": label if label else f"Chat {len(self.conversations_list) + 1}", "key": session_id}
-        self.conversations_list.insert(0, new_conv_entry) # Add to top
-        return self._get_current_history(), self._update_conversation_list_ui()
-    
-    def change_chat_session(self, session_id):
-        if session_id and session_id in self.conversation_contexts:
-            self.current_conversation_id = session_id
-            # Potentially update hyperparameter UI elements based on loaded session_settings
-            # For now, just load history
-            return self._get_current_history()
-        return self._get_current_history() # No change or invalid ID
-
-    def clear_history(self):
-        if self.current_conversation_id:
-            self.conversation_contexts[self.current_conversation_id]["history"] = [DEFAULT_WELCOME_MESSAGE.copy()]
-            # Also clear current inputs if desired
-            return self._get_current_history(), gr.update(value=None), gr.update(value=None)
-        return [], gr.update(value=None), gr.update(value=None)
-
-    def build_ui(self):
-        with gr.Blocks(theme=gr.themes.Soft()) as demo:
-            gr.Markdown("""
+
+# Helper function to load example images
+def load_example_image(image_path):
+    try:
+        return Image.open(image_path)
+    except Exception as e:
+        print(f"Error loading example image: {e}")
+        return None
+
+
+# Gradio UI 
+with gr.Blocks() as demo:
+    gr.Markdown("""
 <div>
   <img src="https://lf3-static.bytednsdoc.com/obj/eden-cn/nuhojubrps/banner.png" alt="BAGEL" width="380"/>
-  <h1>Unified BAGEL Chat Interface</h1>
 </div>
 """)
-            
-            with gr.Row():
-                with gr.Column(scale=1):
-                    gr.Markdown("### Conversations")
-                    conversation_selector = gr.Radio(
-                        label="Select Chat", 
-                        choices=[], 
-                        type="value"
-                    )
-                    new_chat_btn = gr.Button("➕ New Chat")
-                    
-                    gr.Markdown("### Operation Mode")
-                    mode_selector = gr.Radio(
-                        label="Select Mode", 
-                        choices=["Text to Image", "Image Edit", "Image Understanding"], 
-                        value="Text to Image",
-                        interactive=True
-                    )
+
+    with gr.Tab("📝 Text to Image"):
+        txt_input = gr.Textbox(
+            label="Prompt", 
+            value="A female cosplayer portraying an ethereal fairy or elf, wearing a flowing dress made of delicate fabrics in soft, mystical colors like emerald green and silver. She has pointed ears, a gentle, enchanting expression, and her outfit is adorned with sparkling jewels and intricate patterns. The background is a magical forest with glowing plants, mystical creatures, and a serene atmosphere."
+        )
+        
+        with gr.Row():
+            show_thinking = gr.Checkbox(label="Thinking", value=False)
+        
+        # Add hyperparameter controls in an accordion
+        with gr.Accordion("Inference Hyperparameters", open=False):
+            # 参数一排两个布局
+            with gr.Group():
+                with gr.Row():
+                    seed = gr.Slider(minimum=0, maximum=1000000, value=0, step=1, 
+                                   label="Seed", info="0 for random seed, positive for reproducible results")
+                    image_ratio = gr.Dropdown(choices=["1:1", "4:3", "3:4", "16:9", "9:16"], 
+                                                value="1:1", label="Image Ratio", 
+                                                info="The longer size is fixed to 1024")
                     
-                    # --- Hyperparameter Accordions ---
-                    # Visibility will be controlled by mode_selector
-                    with gr.Accordion("Text to Image Settings", open=True, visible=True) as tti_accordion:
-                        tti_show_thinking_cb = gr.Checkbox(label="Show Thinking Process", value=False, interactive=True)
-                        tti_seed_slider = gr.Slider(minimum=0, maximum=1000000, value=0, step=1, label="Seed (0 for random)", interactive=True)
-                        tti_image_ratio_dd = gr.Dropdown(choices=["1:1", "4:3", "3:4", "16:9", "9:16"], value="1:1", label="Image Ratio", interactive=True)
-                        tti_cfg_text_scale_slider = gr.Slider(minimum=1.0, maximum=8.0, value=4.0, step=0.1, label="CFG Text Scale", interactive=True)
-                        tti_cfg_interval_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.1, label="CFG Interval Start", interactive=True)
-                        tti_cfg_renorm_type_dd = gr.Dropdown(choices=["global", "local", "text_channel"], value="global", label="CFG Renorm Type", interactive=True)
-                        tti_cfg_renorm_min_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, label="CFG Renorm Min", interactive=True)
-                        tti_num_timesteps_slider = gr.Slider(minimum=10, maximum=100, value=50, step=5, label="Timesteps", interactive=True)
-                        tti_timestep_shift_slider = gr.Slider(minimum=1.0, maximum=5.0, value=3.0, step=0.5, label="Timestep Shift", interactive=True)
-                        with gr.Group(visible=False) as tti_thinking_params_group:
-                            tti_do_sample_cb = gr.Checkbox(label="Sampling (for thinking)", value=False, interactive=True)
-                            tti_max_think_token_slider = gr.Slider(minimum=64, maximum=4096, value=1024, step=64, label="Max Think Tokens", interactive=True)
-                            tti_text_temp_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, label="Temperature (for thinking)", interactive=True)
-                        tti_show_thinking_cb.change(lambda x: gr.update(visible=x), inputs=[tti_show_thinking_cb], outputs=[tti_thinking_params_group])
-
-                    with gr.Accordion("Image Edit Settings", open=False, visible=False) as edit_accordion:
-                        edit_show_thinking_cb = gr.Checkbox(label="Show Thinking Process", value=False, interactive=True)
-                        edit_seed_slider = gr.Slider(minimum=0, maximum=1000000, value=0, step=1, label="Seed (0 for random)", interactive=True)
-                        edit_cfg_text_scale_slider = gr.Slider(1.0, 8.0, value=4.0, step=0.1, label="CFG Text Scale", interactive=True)
-                        edit_cfg_img_scale_slider = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="CFG Image Scale", interactive=True)
-                        edit_cfg_interval_slider = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="CFG Interval Start", interactive=True)
-                        edit_cfg_renorm_type_dd = gr.Dropdown(["global", "local", "text_channel"], value="text_channel", label="CFG Renorm Type", interactive=True)
-                        edit_cfg_renorm_min_slider = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="CFG Renorm Min", interactive=True)
-                        edit_num_timesteps_slider = gr.Slider(10, 100, value=50, step=5, label="Timesteps", interactive=True)
-                        edit_timestep_shift_slider = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Timestep Shift", interactive=True)
-                        with gr.Group(visible=False) as edit_thinking_params_group:
-                            edit_do_sample_cb = gr.Checkbox(label="Sampling (for thinking)", value=False, interactive=True)
-                            edit_max_think_token_slider = gr.Slider(64, 4096, value=1024, step=64, label="Max Think Tokens", interactive=True)
-                            edit_text_temp_slider = gr.Slider(0.1, 1.0, value=0.3, step=0.1, label="Temperature (for thinking)", interactive=True)
-                        edit_show_thinking_cb.change(lambda x: gr.update(visible=x), inputs=[edit_show_thinking_cb], outputs=[edit_thinking_params_group])
-
-                    with gr.Accordion("Image Understanding Settings", open=False, visible=False) as und_accordion:
-                        und_show_thinking_cb = gr.Checkbox(label="Show Thinking Process (if applicable)", value=False, interactive=True)
-                        und_seed_slider = gr.Slider(minimum=0, maximum=1000000, value=0, step=1, label="Seed (0 for random)", interactive=True)
-                        und_do_sample_cb = gr.Checkbox(label="Sampling", value=False, interactive=True)
-                        und_text_temp_slider = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature", interactive=True)
-                        und_max_new_tokens_slider = gr.Slider(32, 2048, value=512, step=32, label="Max New Tokens", interactive=True)
-
-                    # Logic to show/hide accordions based on mode
-                    def update_accordion_visibility(mode):
-                        return (
-                            gr.update(visible=mode == "Text to Image"),
-                            gr.update(visible=mode == "Image Edit"),
-                            gr.update(visible=mode == "Image Understanding")
-                        )
-                    mode_selector.change(update_accordion_visibility, inputs=[mode_selector], outputs=[tti_accordion, edit_accordion, und_accordion])
-
-                with gr.Column(scale=3):
-                    chatbot_ui = gr.Chatbot(label="BAGEL Chat", value=[DEFAULT_WELCOME_MESSAGE.copy()], bubble_full_width=False, height=600)
-                    with gr.Row():
-                        image_upload_ui = gr.Image(type="pil", label="Upload Image (for Edit/Understand)", sources=['upload'], visible=False, interactive=True)
+                with gr.Row():
+                    cfg_text_scale = gr.Slider(minimum=1.0, maximum=8.0, value=4.0, step=0.1, interactive=True,
+                                             label="CFG Text Scale", info="Controls how strongly the model follows the text prompt (4.0-8.0)")
+                    cfg_interval = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.1, 
+                                           label="CFG Interval", info="Start of CFG application interval (end is fixed at 1.0)")
+                
+                with gr.Row():
+                    cfg_renorm_type = gr.Dropdown(choices=["global", "local", "text_channel"], 
+                                                value="global", label="CFG Renorm Type", 
+                                                info="If the genrated image is blurry, use 'global'")
+                    cfg_renorm_min = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, interactive=True,
+                                             label="CFG Renorm Min", info="1.0 disables CFG-Renorm")
+                
+                with gr.Row():
+                    num_timesteps = gr.Slider(minimum=10, maximum=100, value=50, step=5, interactive=True,
+                                            label="Timesteps", info="Total denoising steps")
+                    timestep_shift = gr.Slider(minimum=1.0, maximum=5.0, value=3.0, step=0.5, interactive=True,
+                                             label="Timestep Shift", info="Higher values for layout, lower for details")
+                
+                # Thinking parameters in a single row
+                thinking_params = gr.Group(visible=False)
+                with thinking_params:
                     with gr.Row():
-                        text_input_ui = gr.Textbox(label="Enter your prompt here...", lines=3, scale=7, interactive=True)
-                        submit_btn = gr.Button("Send", variant="primary", scale=1)
-                        clear_btn = gr.Button("Clear Chat", scale=1)
-                    
-                    # Show/hide image upload based on mode
-                    def update_image_upload_visibility(mode):
-                        return gr.update(visible=mode in ["Image Edit", "Image Understanding"])
-                    mode_selector.change(update_image_upload_visibility, inputs=[mode_selector], outputs=[image_upload_ui])
-
-            # Initial state setup
-            demo.load(lambda: self.new_chat_session("Welcome Chat"), outputs=[chatbot_ui, conversation_selector])
-
-            # Event handlers
-            new_chat_btn.click(
-                self.new_chat_session,
-                inputs=None, 
-                outputs=[chatbot_ui, conversation_selector]
-            )
-            conversation_selector.change(
-                self.change_chat_session,
-                inputs=[conversation_selector],
-                outputs=[chatbot_ui]
+                        do_sample = gr.Checkbox(label="Sampling", value=False, info="Enable sampling for text generation")
+                        max_think_token_n = gr.Slider(minimum=64, maximum=4006, value=1024, step=64, interactive=True,
+                                                    label="Max Think Tokens", info="Maximum number of tokens for thinking")
+                        text_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, interactive=True,
+                                                  label="Temperature", info="Controls randomness in text generation")
+        
+        thinking_output = gr.Textbox(label="Thinking Process", visible=False)
+        img_output = gr.Image(label="Generated Image")
+        gen_btn = gr.Button("Generate")
+        
+        # Dynamically show/hide thinking process box and parameters
+        def update_thinking_visibility(show):
+            return gr.update(visible=show), gr.update(visible=show)
+        
+        show_thinking.change(
+            fn=update_thinking_visibility,
+            inputs=[show_thinking],
+            outputs=[thinking_output, thinking_params]
+        )
+        
+        # Process function based on thinking option and hyperparameters
+        @spaces.GPU(duration=90)
+        def process_text_to_image(prompt, show_thinking, cfg_text_scale, 
+                                 cfg_interval, timestep_shift, 
+                                 num_timesteps, cfg_renorm_min, cfg_renorm_type, 
+                                 max_think_token_n, do_sample, text_temperature, seed, image_ratio):
+            image, thinking = text_to_image(
+                prompt, show_thinking, cfg_text_scale, cfg_interval,
+                timestep_shift, num_timesteps, 
+                cfg_renorm_min, cfg_renorm_type,
+                max_think_token_n, do_sample, text_temperature, seed, image_ratio
             )
-
-            submit_btn.click(
-                self.add_message, 
-                inputs=[
-                    text_input_ui, image_upload_ui, mode_selector, 
-                    # TTI
-                    tti_show_thinking_cb, tti_cfg_text_scale_slider, tti_cfg_interval_slider, tti_timestep_shift_slider, tti_num_timesteps_slider, tti_cfg_renorm_min_slider, tti_cfg_renorm_type_dd, tti_max_think_token_slider, tti_do_sample_cb, tti_text_temp_slider, tti_seed_slider, tti_image_ratio_dd,
-                    # Edit
-                    edit_show_thinking_cb, edit_cfg_text_scale_slider, edit_cfg_img_scale_slider, edit_cfg_interval_slider, edit_timestep_shift_slider, edit_num_timesteps_slider, edit_cfg_renorm_min_slider, edit_cfg_renorm_type_dd, edit_max_think_token_slider, edit_do_sample_cb, edit_text_temp_slider, edit_seed_slider,
-                    # Understand
-                    und_show_thinking_cb, und_do_sample_cb, und_text_temp_slider, und_max_new_tokens_slider, und_seed_slider
-                ],
-                outputs=[chatbot_ui, text_input_ui, image_upload_ui] 
+            return image, thinking if thinking else ""
+        
+        gen_btn.click(
+            fn=process_text_to_image,
+            inputs=[
+                txt_input, show_thinking, cfg_text_scale, 
+                cfg_interval, timestep_shift, 
+                num_timesteps, cfg_renorm_min, cfg_renorm_type,
+                max_think_token_n, do_sample, text_temperature, seed, image_ratio
+            ],
+            outputs=[img_output, thinking_output]
+        )
+
+    with gr.Tab("🖌️ Image Edit"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                edit_image_input = gr.Image(label="Input Image", value=load_example_image('test_images/women.jpg'))
+                edit_prompt = gr.Textbox(
+                    label="Prompt",
+                    value="She boards a modern subway, quietly reading a folded newspaper, wearing the same clothes."
+                )
+            
+            with gr.Column(scale=1):
+                edit_image_output = gr.Image(label="Result")
+                edit_thinking_output = gr.Textbox(label="Thinking Process", visible=False)
+        
+        with gr.Row():
+            edit_show_thinking = gr.Checkbox(label="Thinking", value=False)
+        
+        # Add hyperparameter controls in an accordion
+        with gr.Accordion("Inference Hyperparameters", open=False):
+            with gr.Group():
+                with gr.Row():
+                    edit_seed = gr.Slider(minimum=0, maximum=1000000, value=0, step=1, interactive=True,
+                                        label="Seed", info="0 for random seed, positive for reproducible results")
+                    edit_cfg_text_scale = gr.Slider(minimum=1.0, maximum=8.0, value=4.0, step=0.1, interactive=True,
+                                                  label="CFG Text Scale", info="Controls how strongly the model follows the text prompt")
+                
+                with gr.Row():
+                    edit_cfg_img_scale = gr.Slider(minimum=1.0, maximum=4.0, value=2.0, step=0.1, interactive=True,
+                                                 label="CFG Image Scale", info="Controls how much the model preserves input image details")
+                    edit_cfg_interval = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, interactive=True,
+                                                label="CFG Interval", info="Start of CFG application interval (end is fixed at 1.0)")
+                    
+                with gr.Row():
+                    edit_cfg_renorm_type = gr.Dropdown(choices=["global", "local", "text_channel"], 
+                                                     value="text_channel", label="CFG Renorm Type", 
+                                                     info="If the genrated image is blurry, use 'global")
+                    edit_cfg_renorm_min = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, interactive=True,
+                                                  label="CFG Renorm Min", info="1.0 disables CFG-Renorm")
+                
+                with gr.Row():
+                    edit_num_timesteps = gr.Slider(minimum=10, maximum=100, value=50, step=5, interactive=True,
+                                                 label="Timesteps", info="Total denoising steps")
+                    edit_timestep_shift = gr.Slider(minimum=1.0, maximum=10.0, value=3.0, step=0.5, interactive=True,
+                                                  label="Timestep Shift", info="Higher values for layout, lower for details")
+                
+                
+                # Thinking parameters in a single row
+                edit_thinking_params = gr.Group(visible=False)
+                with edit_thinking_params:
+                    with gr.Row():
+                        edit_do_sample = gr.Checkbox(label="Sampling", value=False, info="Enable sampling for text generation")
+                        edit_max_think_token_n = gr.Slider(minimum=64, maximum=4006, value=1024, step=64, interactive=True,
+                                                         label="Max Think Tokens", info="Maximum number of tokens for thinking")
+                        edit_text_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, interactive=True,
+                                                        label="Temperature", info="Controls randomness in text generation")
+        
+        edit_btn = gr.Button("Submit")
+        
+        # Dynamically show/hide thinking process box for editing
+        def update_edit_thinking_visibility(show):
+            return gr.update(visible=show), gr.update(visible=show)
+        
+        edit_show_thinking.change(
+            fn=update_edit_thinking_visibility,
+            inputs=[edit_show_thinking],
+            outputs=[edit_thinking_output, edit_thinking_params]
+        )
+        
+        # Process editing with thinking option and hyperparameters
+        @spaces.GPU(duration=90)
+        def process_edit_image(image, prompt, show_thinking, cfg_text_scale, 
+                              cfg_img_scale, cfg_interval, 
+                              timestep_shift, num_timesteps, cfg_renorm_min, 
+                              cfg_renorm_type, max_think_token_n, do_sample, 
+                              text_temperature, seed):
+            edited_image, thinking = edit_image(
+                image, prompt, show_thinking, cfg_text_scale, cfg_img_scale, 
+                cfg_interval, timestep_shift, 
+                num_timesteps, cfg_renorm_min, cfg_renorm_type,
+                max_think_token_n, do_sample, text_temperature, seed
             )
-            text_input_ui.submit(
-                 self.add_message, 
-                inputs=[
-                    text_input_ui, image_upload_ui, mode_selector, 
-                    # TTI
-                    tti_show_thinking_cb, tti_cfg_text_scale_slider, tti_cfg_interval_slider, tti_timestep_shift_slider, tti_num_timesteps_slider, tti_cfg_renorm_min_slider, tti_cfg_renorm_type_dd, tti_max_think_token_slider, tti_do_sample_cb, tti_text_temp_slider, tti_seed_slider, tti_image_ratio_dd,
-                    # Edit
-                    edit_show_thinking_cb, edit_cfg_text_scale_slider, edit_cfg_img_scale_slider, edit_cfg_interval_slider, edit_timestep_shift_slider, edit_num_timesteps_slider, edit_cfg_renorm_min_slider, edit_cfg_renorm_type_dd, edit_max_think_token_slider, edit_do_sample_cb, edit_text_temp_slider, edit_seed_slider,
-                    # Understand
-                    und_show_thinking_cb, und_do_sample_cb, und_text_temp_slider, und_max_new_tokens_slider, und_seed_slider
-                ],
-                outputs=[chatbot_ui, text_input_ui, image_upload_ui] 
+            
+            return edited_image, thinking if thinking else ""
+        
+        edit_btn.click(
+            fn=process_edit_image,
+            inputs=[
+                edit_image_input, edit_prompt, edit_show_thinking, 
+                edit_cfg_text_scale, edit_cfg_img_scale, edit_cfg_interval,
+                edit_timestep_shift, edit_num_timesteps, 
+                edit_cfg_renorm_min, edit_cfg_renorm_type,
+                edit_max_think_token_n, edit_do_sample, edit_text_temperature, edit_seed
+            ],
+            outputs=[edit_image_output, edit_thinking_output]
+        )
+
+    with gr.Tab("🖼️ Image Understanding"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                img_input = gr.Image(label="Input Image", value=load_example_image('test_images/meme.jpg'))
+                understand_prompt = gr.Textbox(
+                    label="Prompt", 
+                    value="Can someone explain what's funny about this meme??"
+                )
+            
+            with gr.Column(scale=1):
+                txt_output = gr.Textbox(label="Result", lines=20)
+        
+        with gr.Row():
+            understand_show_thinking = gr.Checkbox(label="Thinking", value=False)
+        
+        # Add hyperparameter controls in an accordion
+        with gr.Accordion("Inference Hyperparameters", open=False):
+            with gr.Row():
+                understand_do_sample = gr.Checkbox(label="Sampling", value=False, info="Enable sampling for text generation")
+                understand_text_temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.05, interactive=True,
+                                                     label="Temperature", info="Controls randomness in text generation (0=deterministic, 1=creative)")
+                understand_max_new_tokens = gr.Slider(minimum=64, maximum=4096, value=512, step=64, interactive=True,
+                                                   label="Max New Tokens", info="Maximum length of generated text, including potential thinking")
+        
+        img_understand_btn = gr.Button("Submit")
+        
+        # Process understanding with thinking option and hyperparameters
+        @spaces.GPU(duration=90)
+        def process_understanding(image, prompt, show_thinking, do_sample, 
+                                 text_temperature, max_new_tokens):
+            result = image_understanding(
+                image, prompt, show_thinking, do_sample, 
+                text_temperature, max_new_tokens
             )
-
-            clear_btn.click(self.clear_history, inputs=None, outputs=[chatbot_ui, text_input_ui, image_upload_ui])
+            return result
         
-        return demo
+        img_understand_btn.click(
+            fn=process_understanding,
+            inputs=[
+                img_input, understand_prompt, understand_show_thinking,
+                understand_do_sample, understand_text_temperature, understand_max_new_tokens
+            ],
+            outputs=txt_output
+        )
+
+    gr.Markdown("""
+<div style="display: flex; justify-content: flex-start; flex-wrap: wrap; gap: 10px;">
+  <a href="https://bagel-ai.org/">
+    <img
+      src="https://img.shields.io/badge/BAGEL-Website-0A66C2?logo=safari&logoColor=white"
+      alt="BAGEL Website"
+    />
+  </a>
+  <a href="https://arxiv.org/abs/2505.14683">
+    <img
+      src="https://img.shields.io/badge/BAGEL-Paper-red?logo=arxiv&logoColor=red"
+      alt="BAGEL Paper on arXiv"
+    />
+  </a>
+  <a href="https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT">
+    <img 
+        src="https://img.shields.io/badge/BAGEL-Hugging%20Face-orange?logo=huggingface&logoColor=yellow" 
+        alt="BAGEL on Hugging Face"
+    />
+  </a>
+  <a href="https://demo.bagel-ai.org/">
+    <img
+      src="https://img.shields.io/badge/BAGEL-Demo-blue?logo=googleplay&logoColor=blue"
+      alt="BAGEL Demo"
+    />
+  </a>
+  <a href="https://discord.gg/Z836xxzy">
+    <img
+      src="https://img.shields.io/badge/BAGEL-Discord-5865F2?logo=discord&logoColor=purple"
+      alt="BAGEL Discord"
+    />
+  </a>
+  <a href="mailto:[email protected]">
+    <img
+      src="https://img.shields.io/badge/BAGEL-Email-D14836?logo=gmail&logoColor=red"
+      alt="BAGEL Email"
+    />
+  </a>
+</div>
+""")
 
-# Main execution
-if __name__ == "__main__":
-    app_instance = GradioApp()
-    demo_ui = app_instance.build_ui()
-    demo_ui.queue().launch(share=True, debug=True) # Set share=True if you need a public link
+demo.launch()