diff --git a/deploy/deploy_server.py b/deploy/deploy_server.py
index fb0e09d..7f633da 100644
--- a/deploy/deploy_server.py
+++ b/deploy/deploy_server.py
@@ -125,7 +125,8 @@ def generate_diffusion_model_api(data: dict):
 
         return JSONResponse({"code": 200, "msg": "success!", "data": json.dumps(re_json, ensure_ascii=False)})
     except Exception as e:
-        return JSONResponse({"code": 500, "msg": str(e), "data": None})
+        logger.exception("Exception occurred while generating diffusion model")
+        return JSONResponse({"code": 500, "msg": "An internal error has occurred.", "data": None})
 
 
 @app.post("/api/generate/sr")
diff --git a/iddm/tools/generate.py b/iddm/tools/generate.py
index 0d6cbb8..034d4aa 100644
--- a/iddm/tools/generate.py
+++ b/iddm/tools/generate.py
@@ -245,7 +245,7 @@ def init_generate_args():
     # Input image size (required)
     # [Warn] Compatible with older versions
     # [Warn] Version <= 1.1.1 need to be equal to model's image size, version > 1.1.1 can set whatever you want
-    parser.add_argument("--image_size", "-i", type=parse_image_size_type, default=64)
+    parser.add_argument("--image_size", "-i", type=check_parse_image_size_type, default=64)
     # Set the use GPU in generate (required)
     parser.add_argument("--use_gpu", type=int, default=0)