Update app.py

app.py

@@ -11,6 +11,8 @@ import struct
 import termios
 from pathlib import Path
 
+LAST_OUTPUT_FILE = "/tmp/noai_last_output.txt"
+
 
 def set_terminal_size(fd, rows=24, cols=80):
     size = struct.pack('HHHH', rows, cols, 0, 0)
@@ -104,10 +106,14 @@ def make_terminal_html(content, cursor=True):
 
 
 def remove_watermark(file_obj, strength, steps):
-    """Generator: yields (terminal_html, path_for_state)."""
+    """Generator: yields terminal HTML only. Writes output path to disk."""
     if file_obj is None:
         raise gr.Error("Please upload an image first.")
 
+    # Clear previous output
+    if os.path.exists(LAST_OUTPUT_FILE):
+        os.remove(LAST_OUTPUT_FILE)
+
     input_path = Path(file_obj)
     suffix = input_path.suffix or ".png"
     tmp_dir = tempfile.mkdtemp()
@@ -140,25 +146,29 @@ def remove_watermark(file_obj, strength, steps):
         if is_done:
             display = escape_html(screen_text)
             if output_path.exists():
+                # Write path to disk
+                with open(LAST_OUTPUT_FILE, "w") as f:
+                    f.write(str(output_path))
                 display += "\n\n<span style='color:#50fa7b;font-weight:bold'>✅ Done! Click 📥 Download below.</span>"
-                yield make_terminal_html(cmd_html + "\n\n" + display, cursor=False), str(output_path)
             else:
                 display += "\n\n<span style='color:#ff5555;font-weight:bold'>❌ Error: output file not created.</span>"
-                yield make_terminal_html(cmd_html + "\n\n" + display, cursor=False), ""
+            yield make_terminal_html(cmd_html + "\n\n" + display, cursor=False)
             return
 
         if now - last_yield > 0.5:
             display = escape_html(screen_text)
-            yield make_terminal_html(cmd_html + "\n\n" + display), ""
+            yield make_terminal_html(cmd_html + "\n\n" + display)
             last_yield = now
 
 
-def do_download(path_str):
-    """Reads path from State, returns file for gr.File."""
-    print(f"[DEBUG] do_download called with path_str='{path_str}'")
-    if path_str and os.path.exists(path_str):
-        return path_str
-    raise gr.Error("File not ready yet. Please wait for processing to complete, then click download.")
+def do_download():
+    """Reads path from disk, returns file."""
+    if os.path.exists(LAST_OUTPUT_FILE):
+        with open(LAST_OUTPUT_FILE) as f:
+            path = f.read().strip()
+        if path and os.path.exists(path):
+            return path
+    raise gr.Error("No file available. Please run watermark removal first.")
 
 
 def check_ai_metadata(file_obj):
@@ -258,23 +268,19 @@ with gr.Blocks(title="noai-watermark") as demo:
             run_btn = gr.Button("🚀 Remove Watermark", variant="primary", size="lg")
             wm_terminal = gr.HTML(value=make_terminal_html("<span style='color:#565f89'>Waiting for input...</span>"))
 
-            # State stores path per-session
-            path_state = gr.State(value="")
-
+            # Download: completely independent, reads from disk
             download_btn = gr.Button("📥 Download Cleaned Image", variant="secondary")
             download_output = gr.File(label="Cleaned Image", interactive=False)
 
-            # Generator → terminal + state
             run_btn.click(
                 fn=remove_watermark,
                 inputs=[upload_wm, strength, steps],
-                outputs=[wm_terminal, path_state],
+                outputs=[wm_terminal],
             )
 
-            # Download: state → file (completely independent event)
             download_btn.click(
                 fn=do_download,
-                inputs=[path_state],
+                inputs=[],
                 outputs=[download_output],
             )