Update app.py

app.py

@@ -29,14 +29,13 @@ else:
     
     os.environ['PLAYWRIGHT_BROWSERS_PATH'] = "/home/user/huggingface/ms-playwright"
     # os.system("playwright install chromium")
-    # result = subprocess.run(
-    #     [sys.executable, "-m", "playwright", "install", "chromium"],
-    #     env={**os.environ},
-    #     check=True,
-    #     stdout=subprocess.PIPE,
-    #     stderr=subprocess.PIPE
-    # )
-    subprocess.run(["playwright", "install"])
+    result = subprocess.run(
+        [sys.executable, "-m", "playwright", "install", "chromium"],
+        env={**os.environ},
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE
+    )
 
 import gradio as gr
 import fitz  # PyMuPDF
@@ -56,6 +55,20 @@ from transformers import AutoProcessor, AutoModelForImageTextToText
 # tts = None
 # voice_style = None
 
+global_stop_thinking = [False]
+global_kill_threads = [False]
+
+def set_stop_thinking():
+    global_stop_thinking[0] = True
+    print(f"[STOP-THINK] set_stop_thinking CALLED! Flag is now: {global_stop_thinking[0]}")
+    return gr.update(value="⚡ Forcing generation...")
+
+def set_kill_threads():
+    global_kill_threads[0] = True
+    print(f"[STOP-THINK] set_kill_threads CALLED! Flag is now: {global_kill_threads[0]}")
+    return gr.update(value="🛑 Stopping...")
+
+
 def extract_pdf_content(pdf_path, max_pages=2):
     """Extract text and images from up to max_pages of a PDF."""
     doc = fitz.open(pdf_path)
@@ -203,9 +216,20 @@ Text:
             padding=True
         ).to("cuda")
         
-        from transformers import TextIteratorStreamer
+        global_stop_thinking[0] = False
+        global_kill_threads[0] = False
+        print(f"[STOP-THINK] Flags RESET. stop_thinking={global_stop_thinking[0]}, kill={global_kill_threads[0]}")
+        
+        from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
         from threading import Thread
 
+        class StopThinkingCriteria(StoppingCriteria):
+            def __call__(self, input_ids, scores, **kwargs):
+                val = global_stop_thinking[0] or global_kill_threads[0]
+                if val:
+                    print(f"[STOP-THINK] Criteria returning True! stop={global_stop_thinking[0]} kill={global_kill_threads[0]}")
+                return val
+
         streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
         generation_kwargs = dict(
             **inputs,
@@ -213,6 +237,7 @@ Text:
             max_new_tokens=2048*16,
             do_sample=True,
             repetition_penalty=repetition_penalty_val,
+            stopping_criteria=StoppingCriteriaList([StopThinkingCriteria()])
         )
 
         if len(images) > 0:
@@ -220,7 +245,22 @@ Text:
         else:
             generation_kwargs.update(dict(temperature=1.0, top_p=0.95, top_k=20, min_p=0.0))
 
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        generation_result = []
+        def generate_and_capture(**kwargs):
+            try:
+                out = model.generate(**kwargs)
+                generation_result.append(out)
+            except Exception as e:
+                import traceback
+                print(f"\n[THREAD1 ERROR] model.generate crashed: {e}")
+                traceback.print_exc()
+            finally:
+                try:
+                    streamer.end()
+                except Exception:
+                    pass
+
+        thread = Thread(target=generate_and_capture, kwargs=generation_kwargs)
         thread.start()
         
         output_text = ""
@@ -228,6 +268,85 @@ Text:
             output_text += new_text
             yield output_text, None
 
+        thread.join()
+
+        if global_kill_threads[0]:
+            yield output_text + "\n\n[Generation completely stopped by user.]", None
+            return
+
+        if global_stop_thinking[0]:
+            global_stop_thinking[0] = False
+            print(f"[STOP-THINK] INJECTION PATH entered. Reset flag to: {global_stop_thinking[0]}")
+            
+            # Inject the closure of thinking and start of JSON
+            injection_text = "\n</think>\n\n```json\n[\n"
+            output_text += injection_text
+            yield output_text, None
+            
+            # Restart generation with updated context
+            generated_ids = generation_result[0]
+            injection_ids = processor.tokenizer(injection_text, return_tensors="pt", add_special_tokens=False).input_ids.to("cuda")
+            new_input_ids = torch.cat([generated_ids, injection_ids], dim=-1)
+            
+            # Update attention mask
+            new_mask = torch.cat([
+                inputs["attention_mask"],
+                torch.ones((1, new_input_ids.shape[1] - inputs["attention_mask"].shape[1]), dtype=inputs["attention_mask"].dtype, device="cuda")
+            ], dim=-1)
+            
+            new_inputs = {
+                "input_ids": new_input_ids,
+                "attention_mask": new_mask
+            }
+            
+            # Carry over only the visual features; discard stale keys like input_token_type or rope_deltas
+            keys_to_keep = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
+            for k in keys_to_keep:
+                if k in inputs:
+                    new_inputs[k] = inputs[k]
+            
+            class KillCriteria(StoppingCriteria):
+                def __call__(self, input_ids, scores, **kwargs):
+                    return global_kill_threads[0]
+
+            new_streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+            new_generation_kwargs = dict(
+                **new_inputs,
+                streamer=new_streamer,
+                max_new_tokens=2048*16,
+                do_sample=True,
+                repetition_penalty=repetition_penalty_val,
+                stopping_criteria=StoppingCriteriaList([KillCriteria()])
+            )
+            
+            if len(images) > 0:
+                new_generation_kwargs.update(dict(temperature=0.6, top_p=0.95, top_k=20, min_p=0.0))
+            else:
+                new_generation_kwargs.update(dict(temperature=1.0, top_p=0.95, top_k=20, min_p=0.0))
+                
+            def thread2_target(**kwargs):
+                try:
+                    model.generate(**kwargs)
+                except Exception as e:
+                    import traceback
+                    print(f"\n[THREAD2 ERROR] model.generate crashed: {e}")
+                    traceback.print_exc()
+                finally:
+                    # Always unblock the streamer to prevent Gradio UI from hanging permanently
+                    try:
+                        new_streamer.end()
+                    except Exception:
+                        pass
+                    
+            thread2 = Thread(target=thread2_target, kwargs=new_generation_kwargs)
+            thread2.start()
+            
+            for new_text in new_streamer:
+                output_text += new_text
+                yield output_text, None
+                
+            thread2.join()
+
         # DEBUG: Log raw output text
         with open("log/debug_vlm_output.txt", "w", encoding="utf-8") as f:
             f.write(output_text)
@@ -890,7 +1009,7 @@ def create_demo():
             with gr.Column(scale=1):
                 # url_input = gr.Textbox(label="Enter a Website URL 🌐", placeholder=r"e.g. https://storykorean.com/stories?level=beginner&story=tiger", value=r"https://storykorean.com/stories?level=beginner&story=tiger")
                 # https://www.bbc.com/korean/articles/c5yz89k5dw0o
-                url_input = gr.Textbox(label="Enter a Website URL 🌐", placeholder=r"e.g. https://www.bbc.com/korean/articles/c5yz89k5dw0o", value=r"https://www.bbc.com/korean/articles/c5yz89k5dw0o")
+                url_input = gr.Textbox(label="Enter a Website URL 🌐", placeholder=r"e.g. https://www.koreanstudyjunkie.com/post/korean-reading-exercise-for-all-levels-beginner-intermediate-advanced", value=r"https://www.koreanstudyjunkie.com/post/korean-reading-exercise-for-all-levels-beginner-intermediate-advanced")
 
                 pdf_input = gr.File(label="Or Upload Book PDF 📚", file_types=[".pdf"], value=example_pdf)
                 
@@ -911,6 +1030,7 @@ def create_demo():
                 
                 with gr.Row():
                     submit_btn = gr.Button("✨ Generate Flashcards ✨", variant="primary")
+                    stop_thinking_btn = gr.Button("⚡ Stop thinking, Generate now", variant="secondary")
                     stop_btn = gr.Button("🛑 Stop Generation", variant="stop")
             
             with gr.Column(scale=2):
@@ -924,13 +1044,19 @@ def create_demo():
         last_source_state = gr.State(None)
         last_korean_words_state = gr.State(None)
 
+        def reset_btn_text():
+            return gr.update(value="⚡ Stop thinking, Generate now"), gr.update(value="🛑 Stop Generation")
+
+        submit_btn.click(fn=reset_btn_text, inputs=None, outputs=[stop_thinking_btn, stop_btn], queue=False)
+
         generate_event = submit_btn.click(
             fn=process_pdf, 
             inputs=[pdf_input, url_input, translit_lang, translit_format, target_lang, max_text_char_input, repetition_penalty_input, last_source_state, last_korean_words_state], 
             outputs=[output_html, last_source_state, last_korean_words_state, stream_box, extracted_text_box, extracted_images_gallery]
         )
         
-        stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[generate_event])
+        stop_thinking_btn.click(fn=set_stop_thinking, inputs=None, outputs=stop_thinking_btn, queue=False)
+        stop_btn.click(fn=set_kill_threads, inputs=None, outputs=stop_btn, queue=False).then(fn=None, inputs=None, outputs=None, cancels=[generate_event])
         
         # Force autoscroll using Custom JS
         stream_box.change(