Update app.py

app.py

@@ -129,26 +129,27 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-def generate_thinking_html(buffer: str) -> str:
+def progress_with_text(text):
     """
-    Return an HTML snippet with a "Thinking..." label, an animated progress bar,
-    and the current buffered text.
+    Returns an HTML snippet that shows an animated progress bar along with the given text.
     """
-    return f'''
-    <div style="display: flex; align-items: center;">
-        <span style="margin-right: 10px; font-weight: bold;">Thinking...</span>
-        <div style="flex: 1; margin-right: 10px; white-space: pre-wrap;">{buffer}</div>
-        <div style="width: 110px; height: 5px; background: #e0e0e0; position: relative; overflow: hidden;">
-            <div style="width: 100%; height: 100%; background: #1890ff; animation: progressAnimation 1.5s linear infinite;"></div>
-        </div>
+    return f"""
+<div style="display: flex; align-items: center;">
+    <span style="margin-right: 10px;">Thinking...</span>
+    <div style="width: 110px; height: 5px; background-color: #ddd; overflow: hidden; position: relative; margin-left: 10px;">
+      <div style="width: 50%; height: 100%; background-color: #1565c0; animation: loading 1.5s linear infinite;"></div>
     </div>
-    <style>
-    @keyframes progressAnimation {{
-        0% {{ transform: translateX(-100%); }}
-        100% {{ transform: translateX(100%); }}
-    }}
-    </style>
-    '''
+</div>
+<div style="margin-top: 10px;">
+{text}
+</div>
+<style>
+@keyframes loading {{
+    0% {{ transform: translateX(-50%); }}
+    100% {{ transform: translateX(100%); }}
+}}
+</style>
+"""
 
 @spaces.GPU(duration=60, enable_queue=True)
 def generate_image_fn(
@@ -255,6 +256,7 @@ def generate(
         conversation = clean_chat_history(chat_history)
         conversation.append({"role": "user", "content": text})
 
+    # For multimodal input with image files
     if files:
         if len(files) > 1:
             images = [load_image(image) for image in files]
@@ -277,14 +279,15 @@ def generate(
         thread.start()
 
         buffer = ""
-        # Initial yield: progress bar with no text yet.
-        yield gr.HTML(generate_thinking_html(buffer))
+        # Yield the initial animated progress bar with no text yet.
+        yield gr.HTML(progress_with_text(""))
         for new_text in streamer:
             buffer += new_text
             buffer = buffer.replace("<|im_end|>", "")
             time.sleep(0.01)
-            yield gr.HTML(generate_thinking_html(buffer))
+            yield gr.HTML(progress_with_text(buffer))
     else:
+        # For text-only conversation
         input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
         if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -305,17 +308,17 @@ def generate(
         t = Thread(target=model.generate, kwargs=generation_kwargs)
         t.start()
 
+        buffer = ""
+        # Yield the initial animated progress bar with no text yet.
+        yield gr.HTML(progress_with_text(""))
         outputs = []
-        # Initial yield: progress bar with no text yet.
-        yield gr.HTML(generate_thinking_html(""))
         for new_text in streamer:
             outputs.append(new_text)
-            current_text = "".join(outputs)
-            yield gr.HTML(generate_thinking_html(current_text))
+            buffer = "".join(outputs)
+            yield gr.HTML(progress_with_text(buffer))
 
         final_response = "".join(outputs)
-        # Final update: yield the final response as plain text.
-        yield final_response
+        yield gr.HTML(progress_with_text(final_response))
 
         # If TTS was requested, convert the final response to speech.
         if is_tts and voice: