Update style_transfer.py

style_transfer.py

@@ -15,6 +15,7 @@ from diffusers import StableDiffusionImg2ImgPipeline
 import time
 import os
 import base64
+import spaces
 from io import BytesIO
 
 class DogStyleTransfer:
@@ -118,6 +119,7 @@ class DogStyleTransfer:
         print(f"xformers available: {self.xformers_available}")
         print("============================\n")
 
+    @spaces.GPU
     def load_model(self, style_name):
         """Load the appropriate model based on style, handling xformers compatibility"""
         
@@ -207,6 +209,7 @@ class DogStyleTransfer:
 
         return image
 
+    @spaces.GPU
     def transform_style(self, image, style_name, strength=0.75, guidance_scale=7.5):
         """
         Transform image to selected style with improved prompts and parameters