Update app.py

app.py

@@ -1,14 +1,64 @@
 import gradio as gr
 import spaces
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+model_name = "dasomaru/gemma-3-4bit-it-demo"
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+# 🚀 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,
+)
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+@spaces.GPU  # 이 함수 실행될 때 GPU 할당됨!
+def chat(user_input):
+    model.to("cuda")  # 함수 안에서 GPU로 이동!
+    
+    messages = [{
+        "role": "user",
+        "content": [{"type": "text", "text": user_input}]
+    }]
+    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=1.0,
+            top_p=0.95,
+            top_k=64,
+            do_sample=True,
+        )
+
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return output_text.split(user_input)[-1].strip()
+
+# Gradio 인터페이스 설정
+demo = gr.Interface(
+    fn=chat,
+    inputs=gr.Textbox(lines=2, placeholder="Type your message..."),
+    outputs=gr.Textbox(lines=10),
+    title="🧠 Gemma-3 4bit (ZeroGPU)",
+    description="This Space uses the ZeroGPU feature. First request might take a few seconds!"
+)
+
+demo.launch()
+
+
+
+# zero = torch.Tensor([0]).cuda()
+# print(zero.device) # <-- 'cpu' 🤔
+
+# @spaces.GPU
+# def greet(n):
+#     print(zero.device) # <-- 'cuda:0' 🤗
+#     return f"Hello {zero + n} Tensor"
+
+# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
+# demo.launch()