Fix launch args and add decompression progress

app.py

@@ -1,4 +1,5 @@
 import base64
+import inspect
 import os
 import shutil
 import tempfile
@@ -262,14 +263,21 @@ def compress_ui(text, context_window, progress=gr.Progress()):
     return b64, stats_text, file_path
 
 
-def decompress_ui(b64_data, file_data, context_window):
+def decompress_ui(b64_data, file_data, context_window, progress=gr.Progress()):
     raw = _get_compressed_bytes(b64_data, file_data)
     model_path = _resolve_default_model_path()
     tokenizer_path = _resolve_default_tokenizer_path()
     requested_strategy = os.getenv("RWKV_STRATEGY", "cpu fp32")
     effective_strategy = _resolve_strategy()
     model, tokenizer = _load_model_and_tokenizer(model_path, tokenizer_path, effective_strategy)
-    text, stats = decompress_bytes(raw, model, tokenizer, context_window=context_window)
+    text, stats = decompress_bytes(
+        raw,
+        model,
+        tokenizer,
+        context_window=context_window,
+        progress=progress,
+        progress_desc="Decompressing",
+    )
     stats_text = _format_decompress_stats(stats, char_count=len(text))
     if effective_strategy != requested_strategy:
         stats_text += "\n- Strategy: cpu fp32 (forced, CUDA unavailable)"
@@ -349,9 +357,16 @@ def build_ui():
 
 
 if __name__ == "__main__":
-    build_ui().queue(max_size=16).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_api=False,
-    )
+    launch_kwargs = {
+        "server_name": "0.0.0.0",
+        "server_port": 7860,
+        "share": False,
+    }
+    try:
+        launch_params = inspect.signature(gr.Blocks.launch).parameters
+        if "show_api" in launch_params:
+            launch_kwargs["show_api"] = False
+    except (TypeError, ValueError):
+        pass
+
+    build_ui().queue(max_size=16).launch(**launch_kwargs)

llm_compressor.py

@@ -289,7 +289,14 @@ def compress_text(text, model, tokenizer, context_window=2048):
     return compress_tokens(tokens, model, context_window=context_window, original_bytes=original_bytes)
 
 
-def decompress_bytes(data, model, tokenizer, context_window=2048):
+def decompress_bytes(
+    data,
+    model,
+    tokenizer,
+    context_window=2048,
+    progress=None,
+    progress_desc="Decompressing",
+):
     if context_window <= 0:
         raise ValueError("context_window must be positive.")
     if not data or len(data) < 4:
@@ -306,9 +313,12 @@ def decompress_bytes(data, model, tokenizer, context_window=2048):
     context_tokens = []
     state = None
     start_time = time.time()
+    if progress is not None:
+        progress((0, total_tokens), desc=progress_desc, unit="token")
+    progress_step = max(1, total_tokens // 100)
 
     with torch.inference_mode():
-        for _ in range(total_tokens):
+        for idx in range(total_tokens):
             if len(context_tokens) >= context_window:
                 context_tokens = []
                 state = None
@@ -334,6 +344,8 @@ def decompress_bytes(data, model, tokenizer, context_window=2048):
             low_val = int(cdf[target_token_id - 1].item()) if target_token_id > 0 else 0
             high_val = int(cdf[target_token_id].item())
             decoder.update_range(low_val, high_val, total_count)
+            if progress is not None and (idx + 1 == total_tokens or (idx + 1) % progress_step == 0):
+                progress((idx + 1, total_tokens), desc=progress_desc, unit="token")
 
     text = decode_tokens(tokenizer, decoded_tokens)
     duration = time.time() - start_time