Update app.py

app.py

@@ -68,30 +68,26 @@ stds = [0.229, 0.224, 0.225]
 
 img_transforms = transforms.Compose([                        
     transforms.ToTensor(),
-    transforms.Normalize(means,stds)
+    transforms.Normalize(means, stds)
 ])
 
-# Load models globally on CPU
-modelv4_cpu = torch.jit.load(modelarcanev4, map_location='cpu').eval()
-modelv3_cpu = torch.jit.load(modelarcanev3, map_location='cpu').eval()
-modelv2_cpu = torch.jit.load(modelarcanev2, map_location='cpu').eval()
-
 @spaces.GPU
 def proc_pil_img(input_image, model_path):
-    """GPU-accelerated image processing"""
-    # Load model fresh on GPU to avoid device mismatch
+    """GPU-accelerated image processing with half precision support"""
+    # Load model on GPU
     model = torch.jit.load(model_path, map_location='cuda').eval()
     
-    # Create tensors on GPU
-    t_stds = torch.tensor(stds).cuda().view(3, 1, 1)
-    t_means = torch.tensor(means).cuda().view(3, 1, 1)
+    # Create tensors on GPU in half precision to match model
+    t_stds = torch.tensor(stds).cuda().half().view(3, 1, 1)
+    t_means = torch.tensor(means).cuda().half().view(3, 1, 1)
     
-    # Transform and move to GPU
-    transformed_image = img_transforms(input_image).unsqueeze(0).cuda()
+    # Transform image and move to GPU with half precision
+    transformed_image = img_transforms(input_image).unsqueeze(0).cuda().half()
     
     with torch.no_grad():
         result_image = model(transformed_image)[0]
-        output_image = result_image.mul(t_stds).add(t_means).mul(255.).clamp(0, 255).permute(1, 2, 0)
+        # Convert back to float for post-processing
+        output_image = result_image.float().mul(t_stds.float()).add(t_means.float()).mul(255.).clamp(0, 255).permute(1, 2, 0)
         output_image = output_image.cpu().numpy().astype('uint8')
         output_image = PIL.Image.fromarray(output_image)