Update app.py

app.py

@@ -172,20 +172,20 @@ def run_inference(input_tensor):
     with torch.no_grad():
         with torch.autocast(device_type=device.split(':')[0], dtype=torch.bfloat16):
             prediction = model(input_batch)
-    
-    # *** FIX: Convert from BFloat16 to Float32 before returning, making it NumPy compatible ***
     return prediction.cpu().to(torch.float32)
 
 def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name):
     if last_input_map is None: return None, None, None
     c_idx = SDO_CHANNELS.index(channel_name)
     
-    scalers_dict = APP_CACHE["scalers"]
-    scaler = scalers_dict[channel_name]
-    params = scaler.to_dict()
-    mean, std, epsilon, sl_scale_factor = params['mean'], params['std'], params['epsilon'], params['sl_scale_factor']
+    # *** FIX: Access the specific scaler for the channel from the dictionary ***
+    scaler = APP_CACHE["scalers"][channel_name]
+    # *** FIX: Access the parameters as attributes, not from to_dict() ***
+    mean = scaler.mean
+    std = scaler.std
+    epsilon = scaler.epsilon
+    sl_scale_factor = scaler.sl_scale_factor
 
-    # The prediction_tensor is now Float32, so .numpy() will work
     pred_slice = inverse_transform_single_channel(
         prediction_tensor[0, c_idx].numpy(), mean=mean, std=std, epsilon=epsilon, sl_scale_factor=sl_scale_factor
     )