Update inf.py

inf.py

@@ -18,15 +18,6 @@ class InferencePipeline:
         self.hf_token = hf_token
         self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
         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)
-        # else:
-        #     self.pipe = StableDiffusionXLPipeline.from_pretrained(
-        #         self.base_model_id,
-        #         torch_dtype=torch.float16,
-        #         use_auth_token=self.hf_token)
-            # self.pipe = self.pipe.to(self.device)
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
                 self.base_model_id,
                 torch_dtype=torch.float16,
@@ -98,15 +89,10 @@ class InferencePipeline:
 
         self.content_lora_model_id = content_lora_model_id
         self.style_lora_model_id = style_lora_model_id
-    
+
     @spaces.GPU
-    def run(
-            self,
-            content_lora_model_id: str,
-            style_lora_model_id: str,
+    def inference(self,
             prompt: str,
-            content_alpha: float,
-            style_alpha: float,
             seed: int,
             n_steps: int,
             guidance_scale: float,
@@ -114,13 +100,8 @@ class InferencePipeline:
     ) -> PIL.Image.Image:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
-
-        self.load_pipe(content_lora_model_id, style_lora_model_id, content_alpha, style_alpha)
-        
         self.pipe.to("cuda")
-
         generator = torch.Generator(device="cuda").manual_seed(seed)
-        print(self.pipe.device)
         out = self.pipe(
             prompt,
             num_inference_steps=n_steps,
@@ -129,3 +110,26 @@ class InferencePipeline:
             num_images_per_prompt=num_images_per_prompt,
         )  # type: ignore
         return out.images
+    
+    
+    def run(
+            self,
+            content_lora_model_id: str,
+            style_lora_model_id: str,
+            prompt: str,
+            content_alpha: float,
+            style_alpha: float,
+            seed: int,
+            n_steps: int,
+            guidance_scale: float,
+            num_images_per_prompt: int = 1
+    ) -> PIL.Image.Image:
+        
+        self.load_pipe(content_lora_model_id, style_lora_model_id, content_alpha, style_alpha)
+        
+        return self.inference(
+            prompt=prompt,
+            n_steps=n_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+        )