Upload prompt_generator.py

prompt_generator.py

@@ -1,9 +1,14 @@
 import os
 import logging
 import torch
+import utils
 from typing import Tuple, Optional, Dict, Any
 import gc
 
+if utils.is_space_environment():
+    import spaces
+
+
 # ロギング設定
 logging.basicConfig(
     level=logging.INFO,
@@ -99,8 +104,7 @@ def load_model():
         # HuggingFaceからモデルを直接ロードする
         model_path = model_name
 
-        is_spaces_environment = "SPACE_ID" in os.environ and os.environ.get("SYSTEM") == "spaces"
-        if is_spaces_environment:
+        if utils.is_space_environment():
             # Spaces環境ではdevice_mapの設定を変更
             device_map = None
             torch_dtype = torch.float16