Upload 2 files

app.py

@@ -64,7 +64,7 @@ vae = None
 # スタイルリストから名前のみを抽出
 style_names = [style["name"] for style in style_list]
 
-@spaces.GPU(timeout_seconds=300)
+@spaces.GPU(timeout_seconds=120)
 def initialize_llm():
     """アプリケーション起動時にLLMだけを初期化する関数"""
 
@@ -270,7 +270,7 @@ def convert_text_to_prompt(
         return f"エラーが発生しました: {str(e)}", novel_text
 
 @spaces.GPU
-def load_image_model(timeout_seconds=300):
+def load_image_model(timeout_seconds=120):
     """画像生成モデルをロードする関数"""
     global pipe, vae
     

prompt_generator.py

@@ -79,6 +79,7 @@ _model = None
 _tokenizer = None
 
 
+@spaces.GPU
 def load_model():
     """モデルをロードする関数"""
     global _model, _tokenizer
@@ -133,8 +134,6 @@ def load_model():
         logger.error(f"Failed to load prompt generation model: {str(e)}")
         raise
 
-
-
 def unload_model():
     """メモリからモデルをアンロードする関数"""
     global _model, _tokenizer
@@ -154,6 +153,7 @@ def unload_model():
     
     logger.info("Prompt generation model unloaded")
 
+@spaces.GPU
 def generate_prompt(
     novel_text: str,
     series_name: str = "original",
@@ -163,6 +163,7 @@ def generate_prompt(
     try:
         # モデルとトークナイザーの読み込み
         model, tokenizer = load_model()
+        logger.info("Loading model, tokenizer is ok...")
         
         # 入力の検証
         if not novel_text or novel_text.isspace():
@@ -264,6 +265,7 @@ masterpiece, best quality, highresなどの品質に関連するタグは後工
             add_generation_prompt=True,
             return_tensors="pt",
         ).to(model.device)
+        logger.info("tokenizer.apply_chat_template is ok...")
         
         # 長すぎる入力のトリミング
         if inputs.shape[1] > max_input_length:
@@ -271,6 +273,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")
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=inputs,
@@ -286,6 +289,7 @@ masterpiece, best quality, highresなどの品質に関連するタグは後工
                 pad_token_id=tokenizer.pad_token_id,
             )
 
+        logger.info("after ttorch.no_grad")
         # デコード
         full_outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)