Removed streaming

app.py

@@ -82,6 +82,10 @@ def _run_stage_gpu(
 
     if target_emb is None and encode_text_input is not None:
         target_emb = encode_text(encode_text_input, encoder)
+    elif target_emb is not None:
+        # Move CPU tensor back to GPU for beam search
+        device = next(llm.parameters()).device
+        target_emb = target_emb.to(device)
 
     step_count = 0
     def count_steps(step, cand):
@@ -101,7 +105,12 @@ def _run_stage_gpu(
         on_step=count_steps,
     )
     elapsed = time.time() - t0
-    return result, elapsed, step_count, target_emb
+    # Return only CPU/plain data to avoid CUDA init in main process on ZeroGPU
+    return {
+        "seq_str": result.seq_str,
+        "cos_sim": result.cos_sim,
+        "token_ids": result.token_ids,
+    }, elapsed, step_count, target_emb.cpu()
 
 
 def run_stage(
@@ -131,20 +140,20 @@ def run_stage(
     # On Stage 1, pass raw text so encoding happens inside GPU context
     encode_input = text.strip() if stage_num == 1 else None
 
-    result, elapsed, steps, returned_emb = _run_stage_gpu(
+    result_dict, elapsed, steps, returned_emb_cpu = _run_stage_gpu(
         target_emb_state, encoder_name, prompt,
         beam_width, top_k, patience, max_steps, min_similarity, randomness,
         encode_text_input=encode_input,
     )
 
-    if returned_emb is not None:
-        target_emb_state = returned_emb
+    # Store embedding on CPU — it gets moved back to GPU inside _run_stage_gpu
+    target_emb_state = returned_emb_cpu
 
     stage_results_state = stage_results_state + [{
         "stage": stage_num,
-        "text": result.seq_str,
-        "cos_sim": result.cos_sim,
-        "length": len(result.token_ids),
+        "text": result_dict["seq_str"],
+        "cos_sim": result_dict["cos_sim"],
+        "length": len(result_dict["token_ids"]),
         "time": elapsed,
         "steps": steps,
     }]