Update style_transfer.py

style_transfer.py

@@ -8,6 +8,9 @@ import os
 import base64
 from io import BytesIO
 
+os.environ['TORCH_ENABLE_MPS_FALLBACK'] = '1'
+os.environ['PYTHONWARNINGS'] = 'ignore::DeprecationWarning,ignore::UserWarning'
+
 class DogStyleTransfer:
     """
     Class for handling dog image style transfer using Stable Diffusion.
@@ -124,7 +127,8 @@ class DogStyleTransfer:
                     model_id,
                     cache_dir=self.model_cache_dir,
                     torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                    safety_checker=None  # Remove safety checker to improve speed
+                    safety_checker=None,  # Remove safety checker to improve speed
+                    revision="main"
                 )
 
                 if self.device == "cuda":