Update app.py

app.py

@@ -36,16 +36,15 @@ class ModelManager:
             print("토크나이저 로딩 완료")
             
             print("모델 로딩 시작...")
+            # CUDA 초기화 방지를 위한 설정
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 token=HF_TOKEN,
                 torch_dtype=torch.float16,
-                device_map="auto",
+                device_map=None,  # 초기에는 device_map을 설정하지 않음
                 trust_remote_code=True,
-                low_cpu_mem_usage=True,
-                max_memory={0: "13GB"}  # GPU 메모리 제한
+                low_cpu_mem_usage=True
             )
-            self.model.eval()
             print("모델 로딩 완료")
             
         except Exception as e:
@@ -55,11 +54,13 @@ class ModelManager:
     @spaces.GPU
     def generate_text(self, prompt, max_tokens, temperature, top_p):
         try:
+            # GPU 컨텍스트 내에서 device 설정
+            self.model = self.model.to("cuda")
             input_ids = self.tokenizer.encode(
                 prompt,
                 return_tensors="pt",
                 add_special_tokens=True
-            ).to(self.model.device)
+            ).to("cuda")
 
             with torch.no_grad():
                 output_ids = self.model.generate(
@@ -73,11 +74,15 @@ class ModelManager:
                     num_return_sequences=1
                 )
 
+            # CPU로 다시 이동
+            self.model = self.model.to("cpu")
             return self.tokenizer.decode(
                 output_ids[0][input_ids.shape[1]:],
                 skip_special_tokens=True
             )
         except Exception as e:
+            if self.model.device.type == "cuda":
+                self.model = self.model.to("cpu")
             raise Exception(f"텍스트 생성 실패: {e}")
 
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):