Upload 3 files

app.py

@@ -48,6 +48,7 @@ logger = logging.getLogger(__name__)
 
 # Constants
 IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
+HF_TOKEN = os.getenv("HF_TOKEN")
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
 
 # PyTorch settings for better performance and determinism

prompt_generator.py

@@ -119,8 +119,8 @@ def load_model():
         _model = AutoModelForCausalLM.from_pretrained(
             model_path,
             torch_dtype=torch_dtype,
-#            device_map=device_map,
-#            use_cache=True,
+            device_map=device_map,
+            use_cache=True,
             low_cpu_mem_usage=True,
         )
 
@@ -277,7 +277,7 @@ masterpiece, best quality, highresなどの品質に関連するタグは後工
             logger.warning(f"Input tokens were too many and have been truncated to {max_input_length}")
         
         # 生成
-        logger.info("before ttorch.no_grad")
+        logger.info("before torch.no_grad")
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=inputs,