Update inf.py

inf.py

@@ -17,8 +17,7 @@ class InferencePipeline:
     def __init__(self, hf_token: str | None = None):
         self.hf_token = hf_token
         self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         if self.device.type == 'cpu':
             self.pipe = StableDiffusionXLPipeline.from_pretrained(
                 self.base_model_id, use_auth_token=self.hf_token, cache_dir='./cache')
@@ -109,6 +108,8 @@ class InferencePipeline:
             guidance_scale: float,
             num_images_per_prompt: int = 1
     ) -> PIL.Image.Image:
+        
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')