Update app.py

app.py

@@ -1,31 +1,34 @@
-import gradio as gr
-import spaces
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-model_name = "dasomaru/gemma-3-4bit-it-demo"
-
-# 🚀 모델과 토크나이저를 서버 시작 시 한번만 불러온다
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",  # ZeroGPU에서는 자동으로 GPU 할당됨
-    trust_remote_code=True,
-)
-
-@spaces.GPU(duration=300)
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        top_k=50,
-        do_sample=True,
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")
-demo.launch()
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "dasomaru/gemma-3-4bit-it-demo"
+
+# 🚀 tokenizer는 CPU에서도 미리 불러올 수 있음
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# 🚀 model은 CPU로만 먼저 올림 (GPU 아직 없음)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,  # 4bit model이니까
+    trust_remote_code=True,
+)
+
+@spaces.GPU(duration=300)
+def generate_response(prompt):
+    # 모델 및 토크나이저 로딩은 함수 내부에서 수행
+    tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
+    model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
+    model.to("cuda")
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7,
+    top_p=0.9,
+    top_k=50,
+    do_sample=True,)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")
+demo.launch()
+
+