Update app.py

app.py

@@ -1,27 +1,46 @@
-from fastapi import FastAPI
+import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from fastapi import FastAPI
+from transformers import AutoTokenizer, pipeline
+import threading
 
-MODEL = "google/gemma-4-E4B"
+MODEL = os.environ.get("MODEL_NAME", "google/gemma-4-E4B")
+# Para Space grátis, defina MODEL_NAME=google/gemma-4-E2B no settings se E4B falhar.
 
-# Carregamento otimista para CPU
-tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    device_map={"": "cpu"},
-    torch_dtype=torch.float32,
-    low_cpu_mem_usage=True,
-    trust_remote_code=True
-)
+tokenizer = None
+generator = None
+_model_lock = threading.Lock()
+_loading = False
 
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+def load_model():
+    global tokenizer, generator, _loading
+    with _model_lock:
+        if tokenizer is not None and generator is not None:
+            return
+        _loading = True
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+            # carregamento em CPU
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL,
+                device_map={"": "cpu"},
+                torch_dtype="float32",
+                low_cpu_mem_usage=True,
+                trust_remote_code=True
+            )
+            generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+        finally:
+            _loading = False
 
 def generate(prompt):
-    out = generator(prompt, max_new_tokens=128, do_sample=False)
+    if generator is None:
+        load_model()
+    # limite de tokens para reduzir uso de memória
+    out = generator(prompt, max_new_tokens=64, do_sample=False)
     return out[0]["generated_text"]
 
-demo = gr.Interface(fn=generate, inputs=gr.Textbox(lines=4, label="Prompt"), outputs="text", title="Gemma-4-E4B (CPU)")
+demo = gr.Interface(fn=generate, inputs=gr.Textbox(lines=4, label="Prompt"), outputs="text", title="Gemma (Space CPU)")
 
 app = FastAPI()
 app = gr.mount_gradio_app(app, demo, path="/")