diff --git a/moss_web_demo_gradio.py b/moss_web_demo_gradio.py
index 32fceb2..6b7e8a1 100644
--- a/moss_web_demo_gradio.py
+++ b/moss_web_demo_gradio.py
@@ -3,11 +3,10 @@ from transformers.generation.utils import logger
 from huggingface_hub import snapshot_download
 import mdtex2html
 import gradio as gr
-import platform
+import argparse
 import warnings
 import torch
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
 
 try:
     from transformers import MossForCausalLM, MossTokenizer
@@ -19,20 +18,35 @@ except (ImportError, ModuleNotFoundError):
 logger.setLevel("ERROR")
 warnings.filterwarnings("ignore")
 
-model_path = "fnlp/moss-moon-003-sft"
-if not os.path.exists(model_path):
-    model_path = snapshot_download(model_path)
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", default="fnlp/moss-moon-003-sft-int4", 
+                    choices=["fnlp/moss-moon-003-sft", 
+                             "fnlp/moss-moon-003-sft-int8", 
+                             "fnlp/moss-moon-003-sft-int4"], type=str)
+parser.add_argument("--gpu", default="0", type=str)
+args = parser.parse_args()
 
-print("Waiting for all devices to be ready, it may take a few minutes...")
-config = MossConfig.from_pretrained(model_path)
-tokenizer = MossTokenizer.from_pretrained(model_path)
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+num_gpus = len(args.gpu.split(","))
 
-with init_empty_weights():
-    raw_model = MossForCausalLM._from_config(config, torch_dtype=torch.float16)
-raw_model.tie_weights()
-model = load_checkpoint_and_dispatch(
-    raw_model, model_path, device_map="auto", no_split_module_classes=["MossBlock"], dtype=torch.float16
-)
+if ('int8' in args.model_name or 'int4' in args.model_name) and num_gpus > 1:
+    raise ValueError("Quantized models do not support model parallel. Please run on a single GPU (e.g., --gpu 0) or use `fnlp/moss-moon-003-sft`")
+
+config = MossConfig.from_pretrained(args.model_name)
+tokenizer = MossTokenizer.from_pretrained(args.model_name)
+
+if num_gpus > 1:
+    if not os.path.exists(args.model_name):
+        args.model_name = snapshot_download(args.model_name)
+    print("Waiting for all devices to be ready, it may take a few minutes...")
+    with init_empty_weights():
+        raw_model = MossForCausalLM._from_config(config, torch_dtype=torch.float16)
+    raw_model.tie_weights()
+    model = load_checkpoint_and_dispatch(
+        raw_model, args.model_name, device_map="auto", no_split_module_classes=["MossBlock"], dtype=torch.float16
+    )
+else: # on a single gpu
+    model = MossForCausalLM.from_pretrained(args.model_name, trust_remote_code=True).half().cuda()
 
 meta_instruction = \
     """You are an AI assistant whose name is MOSS.
diff --git a/moss_web_demo_streamlit.py b/moss_web_demo_streamlit.py
index a611a04..7441041 100644
--- a/moss_web_demo_streamlit.py
+++ b/moss_web_demo_streamlit.py
@@ -24,7 +24,7 @@ args = parser.parse_args()
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
 num_gpus = len(args.gpu.split(","))
 
-if args.model_name in ["fnlp/moss-moon-003-sft-int8", "fnlp/moss-moon-003-sft-int4"] and num_gpus > 1:
+if ('int8' in args.model_name or 'int4' in args.model_name) and num_gpus > 1:
     raise ValueError("Quantized models do not support model parallel. Please run on a single GPU (e.g., --gpu 0) or use `fnlp/moss-moon-003-sft`")
 
 st.set_page_config(
@@ -34,7 +34,7 @@ st.set_page_config(
      initial_sidebar_state="expanded",
  )
 
-st.title(':robot_face: moss-moon-003-sft')
+st.title(':robot_face: {}'.format(args.model_name.split('/')[-1]))
 st.sidebar.header("Parameters")
 temperature = st.sidebar.slider("Temerature", min_value=0.0, max_value=1.0, value=0.7)
 max_length = st.sidebar.slider('Maximum response length', min_value=256, max_value=1024, value=512)