Update app.py

app.py

@@ -17,7 +17,6 @@ hf_hub_download(
 )
 
 # 推論関数
-@spaces.GPU()
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -28,59 +27,79 @@ def respond(
     top_p,
     top_k,
     repeat_penalty,
+    use_gpu: bool = True  # Añadir parámetro para elegir entre GPU y CPU
 ):
     chat_template = MessagesFormatterType.GEMMA_2
 
-    llm = Llama(
-        model_path=f"models/{model}",
-        flash_attn=True,
-        n_gpu_layers=81,
-        n_batch=1024,
-        n_ctx=8192,
-    )
-    provider = LlamaCppPythonProvider(llm)
+    try:
+        # Si no hay GPU, usar CPU
+        if use_gpu:
+            llm = Llama(
+                model_path=f"models/{model}",
+                flash_attn=True,
+                n_gpu_layers=81,
+                n_batch=1024,
+                n_ctx=8192,
+            )
+        else:
+            llm = Llama(
+                model_path=f"models/{model}",
+                flash_attn=False,  # Desactivar el uso de GPU
+                n_batch=1024,
+                n_ctx=8192,
+            )
 
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=f"{system_message}",
-        predefined_messages_formatter_type=chat_template,
-        debug_output=True
-    )
+        provider = LlamaCppPythonProvider(llm)
+
+        agent = LlamaCppAgent(
+            provider,
+            system_prompt=f"{system_message}",
+            predefined_messages_formatter_type=chat_template,
+            debug_output=True
+        )
 
-    settings = provider.get_provider_default_settings()
-    settings.temperature = temperature
-    settings.top_k = top_k
-    settings.top_p = top_p
-    settings.max_tokens = max_tokens
-    settings.repeat_penalty = repeat_penalty
-    settings.stream = True
+        settings = provider.get_provider_default_settings()
+        settings.temperature = temperature
+        settings.top_k = top_k
+        settings.top_p = top_p
+        settings.max_tokens = max_tokens
+        settings.repeat_penalty = repeat_penalty
+        settings.stream = True
 
-    messages = BasicChatHistory()
+        messages = BasicChatHistory()
 
-    for msn in history:
-        user = {
-            'role': Roles.user,
-            'content': msn[0]
-        }
-        assistant = {
-            'role': Roles.assistant,
-            'content': msn[1]
-        }
-        messages.add_message(user)
-        messages.add_message(assistant)
+        for msn in history:
+            user = {
+                'role': Roles.user,
+                'content': msn[0]
+            }
+            assistant = {
+                'role': Roles.assistant,
+                'content': msn[1]
+            }
+            messages.add_message(user)
+            messages.add_message(assistant)
 
-    stream = agent.get_chat_response(
-        message,
-        llm_sampling_settings=settings,
-        chat_history=messages,
-        returns_streaming_generator=True,
-        print_output=False
-    )
+        stream = agent.get_chat_response(
+            message,
+            llm_sampling_settings=settings,
+            chat_history=messages,
+            returns_streaming_generator=True,
+            print_output=False
+        )
 
-    outputs = ""
-    for output in stream:
-        outputs += output
-        yield outputs
+        outputs = ""
+        for output in stream:
+            outputs += output
+            yield outputs
+    except Exception as e:
+        # Si se supera la cuota de GPU, retornar mensaje de error o intentar con CPU
+        if "You have exceeded your GPU quota" in str(e):
+            print("GPU quota exceeded, switching to CPU mode.")
+            yield "Error: Exceeded GPU quota, switching to CPU. Please wait a moment..."
+            return respond(message, history, model, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty, use_gpu=False)
+        else:
+            yield f"An error occurred: {str(e)}"
 
 # Gradioのインターフェースを作成
 def create_interface(model_name, description):
@@ -137,4 +156,4 @@ with demo:
     interface.render()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()