Update src/streamlit_app.py

src/streamlit_app.py

@@ -137,13 +137,13 @@ model_dir = snapshot_download(
     local_dir_use_symlinks=False    # Tránh tạo symlink vào /.cache
 )
 
-print("Model directory:", model_dir)
-
 # Load model using the local path + token
 model = CLIPModel.from_pretrained(
     model_dir,
-    use_auth_token=hf_token
-).to("cpu").eval()
+    use_auth_token=hf_token,
+    device_map="auto",            # Tự động phân phối weights lên CPU/GPU
+    low_cpu_mem_usage=True,       # Giảm RAM khi load
+).eval()
 
 
 index_path   = os.path.join(model_dir, "faiss_index.bin")
@@ -190,8 +190,15 @@ for idx, img_id in enumerate(example):
 
 # Chạy ví dụ với prompt
 st.subheader("Example usage: enter a prompt to retrieve related images")
-prompt_input = st.text_input("Enter a prompt", "a red Apparel dress")
-top_k_input = st.number_input("Enter the number of results (top_k)", min_value=1, max_value=10, value=5)
-
-if st.button(f"Find top {top_k_input} related images"):
-    running(prompt_input, top_k_input)
+with st.form(key="retrieval_form"):
+    prompt_input = st.text_input("Enter a prompt", placeholder="e.g., a red Apparel dress")
+    top_k_input = st.number_input("Enter the number of results (top_k)", min_value=1, max_value=10, value=5, step=1)
+    
+    submitted = st.form_submit_button(label="Find Related Images")
+
+# Khi người dùng nhấn nút Submit
+if submitted:
+    if prompt_input.strip() and top_k_input > 0:
+        running(prompt_input, top_k_input)
+    else:
+        st.warning("Please enter a valid prompt and top_k.")