Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-# os.environ['CUDA_VISIBLE_DEVICES'] = ''
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
 import spaces
 import sys
 import time
@@ -289,12 +289,10 @@ class LlamaAdapter:
             repeat_penalty=repeat_penalty
         )
     
-@spaces.GPU(duration=120)
 def load_model_gpu(model_type, model_path, n_gpu_layers, params):
     llama = LlamaAdapter(model_path, params, n_gpu_layers)
     print(f"{model_type} モデル {model_path} のロードが完了しました。(n_gpu_layers: {n_gpu_layers})")
     return llama
-
     
 class CharacterMaker:
     def __init__(self):
@@ -336,7 +334,7 @@ class CharacterMaker:
 
             try:
                 # 新しいモデルをロード
-                self.llama = load_model_gpu(model_type, model_path, n_gpu_layers, params)
+                self.llama = LlamaAdapter(model_path, params, n_gpu_layers)
                 self.current_model = model_type
                 self.model_loaded.set()
                 print(f"{model_type} モデルをロードしました。モデルパス: {model_path}、GPUレイヤー数: {n_gpu_layers}")
@@ -344,6 +342,17 @@ class CharacterMaker:
                 print(f"{model_type} モデルのロード中にエラーが発生しました: {e}")
                 self.model_loaded.set()
 
+    @spaces.GPU(duration=120)
+    def chat_or_gen(self, text, gen_characters, gen_token_multiplier, instruction, mode):
+        if mode == "chat":
+            return self.generate_response(text)
+        elif mode == "gen":
+            return self.generate_text(text, gen_characters, gen_token_multiplier, instruction)
+
+
+    def generate_text_gen_pre(self, text, gen_characters, gen_token_multiplier, instruction):
+        return self.chat_or_gen(self, text, gen_characters, gen_token_multiplier, instruction, mode="gen")
+
     def generate_response(self, input_str):
         self.load_model('CHAT')
         if not self.model_loaded.wait(timeout=30) or not self.llama:
@@ -470,14 +479,14 @@ def chat_with_character(message, history):
         character_maker.chat_history = [{"role": "user" if i % 2 == 0 else "assistant", "content": msg} for i, msg in enumerate(sum(history, []))]
     else:
         character_maker.history = [{"user": h[0], "assistant": h[1]} for h in history]
-    return character_maker.generate_response(message)
+    return character_maker.chat_or_gen(text=message,gen_characters=None, gen_token_multiplier=None, instruction=None, mode="chat")
 
 def chat_with_character_stream(message, history):
     if character_maker.use_chat_format:
         character_maker.chat_history = [{"role": "user" if i % 2 == 0 else "assistant", "content": msg} for i, msg in enumerate(sum(history, []))]
     else:
         character_maker.history = [{"user": h[0], "assistant": h[1]} for h in history]
-    response = character_maker.generate_response(message)
+    response = character_maker.chat_or_gen(text=message,gen_characters=None, gen_token_multiplier=None, instruction=None, mode="chat")
     for i in range(len(response)):
         time.sleep(0.05)  # 各文字の表示間隔を調整
         yield response[:i+1]
@@ -702,7 +711,7 @@ def build_gradio_interface():
                 generated_output = gr.Textbox(label="生成された文章")
             
                 generate_button.click(
-                    character_maker.generate_text,
+                    character_maker.generate_text_pre,
                     inputs=[gen_input_text, gen_characters, gen_token_multiplier, gen_instruction],
                     outputs=[generated_output]
                 )