Upload app.py

app.py

@@ -477,6 +477,20 @@ def load_pipeline():
         print(f"Error details: {str(e)}")
         raise RuntimeError(f"Failed to load model: {str(e)}")
 
+def _ensure_pipeline_cpu(pipe):
+    """Ensure Chronos pipeline model/tokenizer are on CPU and in eval mode."""
+    try:
+        pipe.model = pipe.model.to(torch.device("cpu"))
+        pipe.model.eval()
+    except Exception:
+        pass
+    try:
+        tok = getattr(pipe, 'tokenizer', None)
+        if tok is not None and hasattr(tok, 'to'):
+            pipe.tokenizer = tok.to(torch.device("cpu"))
+    except Exception:
+        pass
+
 def is_market_open() -> bool:
     """Check if the US stock market is currently open (legacy function for backward compatibility)"""
     return market_status_manager.get_status('US_STOCKS').is_open
@@ -990,190 +1004,26 @@ def make_prediction_enhanced(symbol: str, timeframe: str = "1d", prediction_days
                 
                 # Use predict_quantiles with proper formatting
                 with torch.amp.autocast('cpu'):
-                    # Ensure all inputs are on CPU
-                    context = context.to(device)
-                    
-                    # Ensure context is properly shaped and on CPU
+                    # Ensure inputs and pipeline are on CPU
                     if len(context.shape) == 1:
                         context = context.unsqueeze(0)
                     context = context.to(device)
-                    
-                    # Force all model components to CPU
-                    pipe.model = pipe.model.to(device)
-                    
-                    # Move model to evaluation mode
-                    pipe.model.eval()
-                    
-                    # Move all model parameters and buffers to CPU
-                    for param in pipe.model.parameters():
-                        param.data = param.data.to(device)
-                    for buffer in pipe.model.buffers():
-                        buffer.data = buffer.data.to(device)
-                    
-                    # Move all model submodules to CPU
-                    for module in pipe.model.modules():
-                        if hasattr(module, 'to'):
-                            module.to(device)
-                    
-                    # Move all model attributes to CPU
-                    for name, value in pipe.model.__dict__.items():
-                        if isinstance(value, torch.Tensor):
-                            pipe.model.__dict__[name] = value.to(device)
-                    
-                    # Move all model config tensors to CPU
-                    if hasattr(pipe.model, 'config'):
-                        for key, value in pipe.model.config.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                setattr(pipe.model.config, key, value.to(device))
-                    
-                    # Move all pipeline tensors to CPU
-                    for name, value in pipe.__dict__.items():
-                        if isinstance(value, torch.Tensor):
-                            setattr(pipe, name, value.to(device))
-                    
-                    # Ensure all model states are on CPU
-                    if hasattr(pipe.model, 'state_dict'):
-                        state_dict = pipe.model.state_dict()
-                        for key in state_dict:
-                            if isinstance(state_dict[key], torch.Tensor):
-                                state_dict[key] = state_dict[key].to(device)
-                        pipe.model.load_state_dict(state_dict)
-                    
-                    # Move any additional components to CPU
-                    if hasattr(pipe, 'tokenizer'):
-                        # Move tokenizer to CPU if it supports it
-                        if hasattr(pipe.tokenizer, 'to'):
-                            pipe.tokenizer = pipe.tokenizer.to(device)
-                        
-                        # Move all tokenizer tensors to CPU
-                        for name, value in pipe.tokenizer.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                setattr(pipe.tokenizer, name, value.to(device))
-                        
-                        # Handle MeanScaleUniformBins specific attributes
-                        if hasattr(pipe.tokenizer, 'bins'):
-                            if isinstance(pipe.tokenizer.bins, torch.Tensor):
-                                pipe.tokenizer.bins = pipe.tokenizer.bins.to(device)
-                        
-                        if hasattr(pipe.tokenizer, 'scale'):
-                            if isinstance(pipe.tokenizer.scale, torch.Tensor):
-                                pipe.tokenizer.scale = pipe.tokenizer.scale.to(device)
-                        
-                        if hasattr(pipe.tokenizer, 'mean'):
-                            if isinstance(pipe.tokenizer.mean, torch.Tensor):
-                                pipe.tokenizer.mean = pipe.tokenizer.mean.to(device)
-                        
-                        # Move any additional tensors in the tokenizer's attributes to CPU
-                        for name, value in pipe.tokenizer.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                pipe.tokenizer.__dict__[name] = value.to(device)
-                        
-                        # Remove the EOS token handling since MeanScaleUniformBins doesn't use it
-                        if hasattr(pipe.tokenizer, '_append_eos_token'):
-                            # Create a wrapper that just returns the input tensors
-                            def wrapped_append_eos(token_ids, attention_mask):
-                                return token_ids, attention_mask
-                            pipe.tokenizer._append_eos_token = wrapped_append_eos
-                    
-                    # Force synchronization for CPU (no-op for CPU)
-                    # torch.cuda.synchronize()  # Not needed for CPU
-                    
-                    # Ensure all model components are in eval mode
-                    pipe.model.eval()
-                    
+                    _ensure_pipeline_cpu(pipe)
+
                     # Fix generation configuration to prevent min_length errors
                     if hasattr(pipe.model, 'config'):
-                        # Ensure generation config is properly set
-                        if hasattr(pipe.model.config, 'generation_config'):
-                            # Reset generation config to safe defaults
-                            pipe.model.config.generation_config.min_length = 0
-                            pipe.model.config.generation_config.max_length = 512
-                            pipe.model.config.generation_config.do_sample = False
-                            pipe.model.config.generation_config.num_beams = 1
-                        else:
-                            # Create a safe generation config if it doesn't exist
+                        gen_cfg = getattr(pipe.model.config, 'generation_config', None)
+                        if gen_cfg is None:
                             pipe.model.config.generation_config = GenerationConfig(
-                                min_length=0,
-                                max_length=512,
-                                do_sample=False,
-                                num_beams=1
+                                min_length=0, max_length=512, do_sample=False, num_beams=1
                             )
-                    
-                    # Move any additional tensors in the model's config to CPU
-                    if hasattr(pipe.model, 'config'):
-                        for key, value in pipe.model.config.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                setattr(pipe.model.config, key, value.to(device))
-                    
-                    # Move any additional tensors in the model's state dict to CPU
-                    if hasattr(pipe.model, 'state_dict'):
-                        state_dict = pipe.model.state_dict()
-                        for key in state_dict:
-                            if isinstance(state_dict[key], torch.Tensor):
-                                state_dict[key] = state_dict[key].to(device)
-                        pipe.model.load_state_dict(state_dict)
-                    
-                    # Move any additional tensors in the model's buffers to CPU
-                    for name, buffer in pipe.model.named_buffers():
-                        if buffer is not None:
-                            pipe.model.register_buffer(name, buffer.to(device))
-                    
-                    # Move any additional tensors in the model's parameters to CPU
-                    for name, param in pipe.model.named_parameters():
-                        if param is not None:
-                            param.data = param.data.to(device)
-                    
-                    # Move any additional tensors in the model's attributes to GPU
-                    for name, value in pipe.model.__dict__.items():
-                        if isinstance(value, torch.Tensor):
-                            pipe.model.__dict__[name] = value.to(device)
-                    
-                    # Move any additional tensors in the model's modules to GPU
-                    for name, module in pipe.model.named_modules():
-                        if hasattr(module, 'to'):
-                            module.to(device)
-                        # Move any tensors in the module's __dict__
-                        for key, value in module.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                setattr(module, key, value.to(device))
-                    
-                    # Force synchronization for CPU (no-op for CPU)
-                    # torch.cuda.synchronize()  # Not needed for CPU
-                    
-                    # Ensure tokenizer is on GPU and all its tensors are on GPU
-                    if hasattr(pipe, 'tokenizer'):
-                        # Move tokenizer to CPU if it supports it
-                        if hasattr(pipe.tokenizer, 'to'):
-                            pipe.tokenizer = pipe.tokenizer.to(device)
-                        
-                        # Move all tokenizer tensors to CPU
-                        for name, value in pipe.tokenizer.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                setattr(pipe.tokenizer, name, value.to(device))
-                        
-                        # Handle MeanScaleUniformBins specific attributes
-                        if hasattr(pipe.tokenizer, 'bins'):
-                            if isinstance(pipe.tokenizer.bins, torch.Tensor):
-                                pipe.tokenizer.bins = pipe.tokenizer.bins.to(device)
-                        
-                        if hasattr(pipe.tokenizer, 'scale'):
-                            if isinstance(pipe.tokenizer.scale, torch.Tensor):
-                                pipe.tokenizer.scale = pipe.tokenizer.scale.to(device)
-                        
-                        if hasattr(pipe.tokenizer, 'mean'):
-                            if isinstance(pipe.tokenizer.mean, torch.Tensor):
-                                pipe.tokenizer.mean = pipe.tokenizer.mean.to(device)
-                        
-                        # Move any additional tensors in the tokenizer's attributes to CPU
-                        for name, value in pipe.tokenizer.__dict__.items():
-                            if isinstance(value, torch.Tensor):
-                                pipe.tokenizer.__dict__[name] = value.to(device)
-                    
-                    # Force synchronization for CPU (no-op for CPU)
-                    # torch.cuda.synchronize()  # Not needed for CPU
-                    
+                        else:
+                            gen_cfg.min_length = 0
+                            gen_cfg.max_length = 512
+                            gen_cfg.do_sample = False
+                            gen_cfg.num_beams = 1
+
                     # Make prediction with proper parameters
-                    # Use the standard quantile levels as per Chronos documentation
                     try:
                         quantiles, mean = pipe.predict_quantiles(
                             context=context,