Add user-configurable GPU duration slider (60-1800 seconds)

- Add GPU Duration slider in UI (default: 600 seconds)
- Refactor to use spaces.GPU context manager with dynamic duration
- Allow users to set GPU time allocation per request
- Maintain backward compatibility with default 600s wrapper

app.py

@@ -269,8 +269,7 @@ def format_validation_message(ok: bool, issues: List[str]) -> str:
     return f"❌ Issues detected:\n{bullets}"
 
 
-@spaces.GPU(duration=600)
-def generate_router_plan_streaming(
+def _generate_router_plan_streaming_internal(
     user_task: str,
     context: str,
     acceptance: str,
@@ -281,8 +280,9 @@ def generate_router_plan_streaming(
     max_new_tokens: int,
     temperature: float,
     top_p: float,
+    gpu_duration: int,
 ):
-    """Generator function for streaming token output."""
+    """Internal generator function for streaming token output."""
     if not user_task.strip():
         yield "", {}, "❌ User task is required.", ""
         return
@@ -291,100 +291,124 @@ def generate_router_plan_streaming(
         yield "", {}, f"❌ Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}", ""
         return
 
-    try:
-        prompt = build_router_prompt(
-            user_task=user_task,
-            context=context,
-            acceptance=acceptance,
-            extra_guidance=extra_guidance,
-            difficulty=difficulty,
-            tags=tags,
-        )
-
-        generator = load_pipeline(model_choice)
-        
-        # Get the underlying model and tokenizer
-        model = generator.model
-        tokenizer = generator.tokenizer
-        
-        # Set up streaming
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        
-        # Prepare inputs
-        inputs = tokenizer(prompt, return_tensors="pt")
-        if hasattr(model, 'device'):
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
-        elif torch.cuda.is_available():
-            inputs = {k: v.cuda() for k, v in inputs.items()}
-        
-        # Start generation in a separate thread
-        generation_kwargs = {
-            **inputs,
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "do_sample": True,
-            "streamer": streamer,
-            "eos_token_id": tokenizer.eos_token_id,
-            "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
-        }
-
-        def _generate():
-            with torch.inference_mode():
-                model.generate(**generation_kwargs)
-
-        thread = Thread(target=_generate)
-        thread.start()
-        
-        # Stream tokens
-        completion = ""
-        parsed_plan: Dict[str, Any] | None = None
-        validation_msg = "🔄 Generating..."
-
-        for new_text in streamer:
-            completion += new_text
-            chunk = completion
-            finished = False
-            display_plan = parsed_plan or {}
-
-            chunk, finished = trim_at_stop_sequences(chunk)
+    # Use GPU context manager with user-specified duration
+    with spaces.GPU(duration=gpu_duration):
+        try:
+            prompt = build_router_prompt(
+                user_task=user_task,
+                context=context,
+                acceptance=acceptance,
+                extra_guidance=extra_guidance,
+                difficulty=difficulty,
+                tags=tags,
+            )
 
-            try:
-                json_block = extract_json_from_text(chunk)
-                candidate_plan = json.loads(json_block)
-                ok, issues = validate_router_plan(candidate_plan)
-                validation_msg = format_validation_message(ok, issues)
-                parsed_plan = candidate_plan if ok else parsed_plan
-                display_plan = candidate_plan
-            except Exception:
-                # Ignore until JSON is complete
-                pass
-
-            yield chunk, display_plan, validation_msg, prompt
-
-            if finished:
-                completion = chunk
-                break
-
-        # Final processing after streaming completes
-        thread.join()
-
-        completion = trim_at_stop_sequences(completion.strip())[0]
-        if parsed_plan is None:
-            try:
-                json_block = extract_json_from_text(completion)
-                parsed_plan = json.loads(json_block)
-                ok, issues = validate_router_plan(parsed_plan)
-                validation_msg = format_validation_message(ok, issues)
-            except Exception as exc:
-                parsed_plan = {}
-                validation_msg = f"❌ JSON parsing failed: {exc}"
-
-        yield completion, parsed_plan, validation_msg, prompt
-        
-    except Exception as exc:
-        error_msg = f"❌ Generation failed: {str(exc)}"
-        yield "", {}, error_msg, ""
+            generator = load_pipeline(model_choice)
+            
+            # Get the underlying model and tokenizer
+            model = generator.model
+            tokenizer = generator.tokenizer
+            
+            # Set up streaming
+            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            
+            # Prepare inputs
+            inputs = tokenizer(prompt, return_tensors="pt")
+            if hasattr(model, 'device'):
+                inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            elif torch.cuda.is_available():
+                inputs = {k: v.cuda() for k, v in inputs.items()}
+            
+            # Start generation in a separate thread
+            generation_kwargs = {
+                **inputs,
+                "max_new_tokens": max_new_tokens,
+                "temperature": temperature,
+                "top_p": top_p,
+                "do_sample": True,
+                "streamer": streamer,
+                "eos_token_id": tokenizer.eos_token_id,
+                "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
+            }
+
+            def _generate():
+                with torch.inference_mode():
+                    model.generate(**generation_kwargs)
+
+            thread = Thread(target=_generate)
+            thread.start()
+            
+            # Stream tokens
+            completion = ""
+            parsed_plan: Dict[str, Any] | None = None
+            validation_msg = "🔄 Generating..."
+
+            for new_text in streamer:
+                completion += new_text
+                chunk = completion
+                finished = False
+                display_plan = parsed_plan or {}
+
+                chunk, finished = trim_at_stop_sequences(chunk)
+
+                try:
+                    json_block = extract_json_from_text(chunk)
+                    candidate_plan = json.loads(json_block)
+                    ok, issues = validate_router_plan(candidate_plan)
+                    validation_msg = format_validation_message(ok, issues)
+                    parsed_plan = candidate_plan if ok else parsed_plan
+                    display_plan = candidate_plan
+                except Exception:
+                    # Ignore until JSON is complete
+                    pass
+
+                yield chunk, display_plan, validation_msg, prompt
+
+                if finished:
+                    completion = chunk
+                    break
+
+            # Final processing after streaming completes
+            thread.join()
+
+            completion = trim_at_stop_sequences(completion.strip())[0]
+            if parsed_plan is None:
+                try:
+                    json_block = extract_json_from_text(completion)
+                    parsed_plan = json.loads(json_block)
+                    ok, issues = validate_router_plan(parsed_plan)
+                    validation_msg = format_validation_message(ok, issues)
+                except Exception as exc:
+                    parsed_plan = {}
+                    validation_msg = f"❌ JSON parsing failed: {exc}"
+
+            yield completion, parsed_plan, validation_msg, prompt
+            
+        except Exception as exc:
+            error_msg = f"❌ Generation failed: {str(exc)}"
+            yield "", {}, error_msg, ""
+
+
+@spaces.GPU(duration=600)  # Default wrapper for backward compatibility
+def generate_router_plan_streaming(
+    user_task: str,
+    context: str,
+    acceptance: str,
+    extra_guidance: str,
+    difficulty: str,
+    tags: str,
+    model_choice: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    gpu_duration: int = 600,
+):
+    """Wrapper function that calls internal generator with GPU duration."""
+    yield from _generate_router_plan_streaming_internal(
+        user_task, context, acceptance, extra_guidance,
+        difficulty, tags, model_choice, max_new_tokens,
+        temperature, top_p, gpu_duration
+    )
 
 
 def clear_outputs():
@@ -446,6 +470,7 @@ def build_ui():
                 max_new_tokens = gr.Slider(256, 20000, value=16000, step=32, label="Max New Tokens")
                 temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature")
                 top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+                gpu_duration = gr.Slider(60, 1800, value=600, step=60, label="GPU Duration (seconds)", info="Maximum GPU time allocation for this request")
 
         generate_btn = gr.Button("Generate Router Plan", variant="primary")
         clear_btn = gr.Button("Clear", variant="secondary")
@@ -469,6 +494,7 @@ def build_ui():
                 max_new_tokens,
                 temperature,
                 top_p,
+                gpu_duration,
             ],
             outputs=[raw_output, plan_json, validation_msg, prompt_view],
             show_progress="full",