Update app.py

app.py

@@ -13,6 +13,7 @@ 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
@@ -31,16 +32,19 @@ save_dir = "./model"
 repo_id = "ByteDance-Seed/BAGEL-7B-MoT"
 cache_dir = save_dir + "/cache"
 
-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"],
-)
+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}")
 
-# Model Initialization
-model_path = "./model" #Download from https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT
+model_path = "./model"
 
 llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
 llm_config.qk_norm = True
@@ -56,7 +60,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,
@@ -77,7 +81,6 @@ 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())},
@@ -97,16 +100,20 @@ 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:
-        if k in device_map:
-            device_map[k] = first_device
-        else:
-            device_map[k] = "cuda:0"
+        device_map[k] = first_device
 else:
-    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
-            
+    # 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
+
 model = load_checkpoint_and_dispatch(
     model,
     checkpoint=os.path.join(model_path, "ema.safetensors"),
@@ -116,8 +123,6 @@ model = load_checkpoint_and_dispatch(
     force_hooks=True,
 ).eval()
 
-
-# Inferencer Preparing 
 inferencer = InterleaveInferencer(
     model=model,
     vae_model=vae_model,
@@ -128,8 +133,7 @@ inferencer = InterleaveInferencer(
 )
 
 def set_seed(seed):
-    """Set random seeds for reproducibility"""
-    if seed > 0:
+    if seed is not None and seed > 0:
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
@@ -140,389 +144,372 @@ def set_seed(seed):
         torch.backends.cudnn.benchmark = False
     return seed
 
-# 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
+# --- 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):
     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],  # End fixed at 1.0
+        cfg_interval=[cfg_interval, 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["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)
+    return result.get("image", None), result.get("text", None) # text is thinking
 
+@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, # Set max_length
+        max_think_token_n=max_new_tokens, 
     )
-    
-    # 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.", ""
-
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
+    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):
+    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)
-    
-    # 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],  # End fixed at 1.0
+        cfg_interval=[cfg_interval, 1.0],
         timestep_shift=timestep_shift,
         num_timesteps=num_timesteps,
         cfg_renorm_min=cfg_renorm_min,
         cfg_renorm_type=cfg_renorm_type,
     )
-    
-    # Include thinking parameter based on user choice
     result = inferencer(image=image, text=prompt, think=show_thinking, **inference_hyper)
-    return result["image"], result.get("text", "")
-
+    return result.get("image", None), result.get("text", None) # text is thinking
+
+# --- Gradio UI --- 
+
+DEFAULT_WELCOME_MESSAGE = {
+    "role": "assistant",
+    "content": [
+        {"type": "text", "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,
+            "image_input_path": image_input.name if image_input else None, # Store path if image is uploaded
+            # 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_message_content = []
+        if text_input:
+            user_message_content.append({"type": "text", "content": text_input})
+        if image_input and mode in ["Image Edit", "Image Understanding"]:
+            # Gradio chatbot can display images directly if they are file paths or PIL Images
+            # For simplicity, let's assume image_input is a PIL Image or path that gr.Image can handle
+            user_message_content.append({"type": "image", "content": image_input})
+        
+        if not user_message_content:
+             user_message_content.append({"type": "text", "content": "(No text prompt provided for image operation)"})
 
-# 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
 
+        history.append({"role": "user", "content": user_message_content, "key": str(uuid.uuid4())})
+        history.append({"role": "assistant", "content": [{"type": "text", "content": "Processing..."}], "key": str(uuid.uuid4()), "loading": True})
+        
+        yield history, gr.update(value=None), gr.update(value=None) # chatbot, text_input, image_input (clear inputs)
 
-# Gradio UI 
-with gr.Blocks() as demo:
-    gr.Markdown("""
+        # Call backend
+        try:
+            output_image = None
+            output_text = None
+            thinking_text = None
+            
+            pil_image_input = Image.open(image_input.name) if image_input else None
+
+            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:
+                bot_response_content.append({"type": "text", "content": f"**Thinking Process:**\n{thinking_text}"}) 
+            if output_text:
+                bot_response_content.append({"type": "text", "content": output_text})
+            if output_image:
+                bot_response_content.append({"type": "image", "content": output_image})
+            
+            if not bot_response_content:
+                bot_response_content.append({"type": "text", "content": "(No output generated)"})
+
+            history[-1]["content"] = bot_response_content
+            history[-1]["loading"] = False
+
+        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("""
 <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.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")
-                    
-                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():
-                        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
-            )
-            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():
+                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")
                     
-                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:
+                    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
+                    )
+                    
+                    # --- 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():
-                        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
+                        image_upload_ui = gr.Image(type="pil", label="Upload Image (for Edit/Understand)", sources=['upload'], visible=False, interactive=True)
+                    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]
             )
-            
-            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
+            conversation_selector.change(
+                self.change_chat_session,
+                inputs=[conversation_selector],
+                outputs=[chatbot_ui]
+            )
+
+            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] 
+            )
+            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 result
+
+            clear_btn.click(self.clear_history, inputs=None, outputs=[chatbot_ui, text_input_ui, image_upload_ui])
         
-        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>
-""")
+        return demo
 
-demo.launch()
+# 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