Update app.py from anycoder

app.py

@@ -3,12 +3,12 @@ import torch
 import gradio as gr
 import spaces
 
-# Load models
+# Load models on CPU initially
 model2 = torch.hub.load(
     "AK391/animegan2-pytorch:main",
     "generator",
     pretrained=True,
-    device="cpu",  # Load on CPU initially
+    device="cpu",
     progress=False
 )
 
@@ -16,29 +16,32 @@ model1 = torch.hub.load(
     "AK391/animegan2-pytorch:main", 
     "generator", 
     pretrained="face_paint_512_v1",  
-    device="cpu"  # Load on CPU initially
+    device="cpu"
 )
 
 face2paint = torch.hub.load(
     'AK391/animegan2-pytorch:main', 
     'face2paint', 
     size=512, 
-    device="cpu",  # Load on CPU initially
+    device="cpu",
     side_by_side=False
 )
 
 @spaces.GPU  # Zero GPU decorator - moves to GPU only when function runs
 def inference(img, ver):
     """Convert portrait to anime style"""
-    # Move models to GPU when function is called
+    # Select model based on version
     if ver == 'Version 2':
-        model2.to('cuda')
-        out = face2paint(model2, img)
-        model2.to('cpu')  # Move back to CPU to free GPU
+        model = model2
     else:
-        model1.to('cuda')
-        out = face2paint(model1, img)
-        model1.to('cpu')  # Move back to CPU to free GPU
+        model = model1
+    
+    # Move model to GPU (spaces.GPU handles the device)
+    model = model.to('cuda')
+    
+    # Process image - face2paint will handle the device internally
+    out = face2paint(model, img, device='cuda')
+    
     return out
 
 # Custom CSS for modern, mobile-friendly design