Update model.py

model.py

@@ -57,7 +57,7 @@ class Model:
             beta_schedule="scaled_linear",
             num_train_timesteps=1000,
             steps_offset=1
-        )
+        ).to('cuda')
         # pipe.enable_freeu(b1=1.1, b2=1.1, s1=0.5, s2=0.7)
         pipe.enable_xformers_memory_efficient_attention()
         pipe.force_zeros_for_empty_prompt = False
@@ -70,34 +70,34 @@ class Model:
         print(f'Loaded {model_id}...')
         return pipe
 
-    def set_base_model(self, base_model_id: str) -> str:
-        if not base_model_id or base_model_id == self.base_model_id:
-            return self.base_model_id
-        del self.pipe
-        torch.cuda.empty_cache()
-        gc.collect()
-        try:
-            self.pipe = self.load_pipe(base_model_id, self.task_name)
-        except Exception:
-            self.pipe = self.load_pipe(self.base_model_id, self.task_name)
-        return self.base_model_id
+    # def set_base_model(self, base_model_id: str) -> str:
+    #     if not base_model_id or base_model_id == self.base_model_id:
+    #         return self.base_model_id
+    #     del self.pipe
+    #     torch.cuda.empty_cache()
+    #     gc.collect()
+    #     try:
+    #         self.pipe = self.load_pipe(base_model_id, self.task_name)
+    #     except Exception:
+    #         self.pipe = self.load_pipe(self.base_model_id, self.task_name)
+    #     return self.base_model_id
 
     def load_controlnet_weight(self, task_name: str) -> None:
         print('Entered load_controlnet_weight....')
-        if task_name == self.task_name:
-            return
-        if self.pipe is not None and hasattr(self.pipe, "controlnet"):
-            del self.pipe.controlnet
-        torch.cuda.empty_cache()
-        gc.collect()
-        model_id = CONTROLNET_MODEL_IDS[task_name]
-        controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
-        print(f'Loaded {model_id}...')
-        controlnet.to(self.device)
-        torch.cuda.empty_cache()
-        gc.collect()
-        self.pipe.controlnet = controlnet
-        self.task_name = task_name
+        # if task_name == self.task_name:
+        #     return
+        # if self.pipe is not None and hasattr(self.pipe, "controlnet"):
+        #     del self.pipe.controlnet
+        # torch.cuda.empty_cache()
+        # gc.collect()
+        # model_id = CONTROLNET_MODEL_IDS[task_name]
+        # controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
+        # print(f'Loaded {model_id}...')
+        # controlnet.to(self.device)
+        # torch.cuda.empty_cache()
+        # gc.collect()
+        # self.pipe.controlnet = controlnet
+        # self.task_name = task_name
 
     def get_prompt(self, prompt: str, additional_prompt: str) -> str:
         if not prompt: