Update app.py

app.py

@@ -129,6 +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:
+    """
+    Return an HTML snippet with a "Thinking..." label, an animated progress bar,
+    and the current buffered 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>
+    </div>
+    <style>
+    @keyframes progressAnimation {{
+        0% {{ transform: translateX(-100%); }}
+        100% {{ transform: translateX(100%); }}
+    }}
+    </style>
+    '''
+
 @spaces.GPU(duration=60, enable_queue=True)
 def generate_image_fn(
     prompt: str,
@@ -256,14 +277,14 @@ def generate(
         thread.start()
 
         buffer = ""
-        yield "Thinking..."
+        # Initial yield: progress bar with no text yet.
+        yield gr.HTML(generate_thinking_html(buffer))
         for new_text in streamer:
             buffer += new_text
             buffer = buffer.replace("<|im_end|>", "")
             time.sleep(0.01)
-            yield buffer
+            yield gr.HTML(generate_thinking_html(buffer))
     else:
-        
         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:]
@@ -285,11 +306,15 @@ def generate(
         t.start()
 
         outputs = []
+        # Initial yield: progress bar with no text yet.
+        yield gr.HTML(generate_thinking_html(""))
         for new_text in streamer:
             outputs.append(new_text)
-            yield "".join(outputs)
+            current_text = "".join(outputs)
+            yield gr.HTML(generate_thinking_html(current_text))
 
         final_response = "".join(outputs)
+        # Final update: yield the final response as plain text.
         yield final_response
 
         # If TTS was requested, convert the final response to speech.
@@ -313,7 +338,6 @@ demo = gr.ChatInterface(
         ["@image Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic"],
         ["Write a Python function to check if a number is prime."],
         ["@tts2 What causes rainbows to form?"],
-
     ],
     cache_examples=False,
     type="messages",
@@ -326,4 +350,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(share=True)
+    demo.queue(max_size=20).launch(share=True)