Update style_transfer.py

style_transfer.py

@@ -24,7 +24,7 @@ class DogStyleTransfer:
     """
     def __init__(self):
         self.models = {}
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = cpu
 
         # Check xformers availability
         self.xformers_available = False
@@ -120,6 +120,11 @@ class DogStyleTransfer:
 
     def load_model(self, style_name):
         """Load the appropriate model based on style, handling xformers compatibility"""
+        
+        if not hasattr(self, '_cuda_initialized'):
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._cuda_initialized = True
+            
         # Get model ID for the style
         model_id = self.style_model_mapping.get(style_name, "runwayml/stable-diffusion-v1-5")