diff --git a/onethingai.py b/onethingai.py
index 92d438f..8258f96 100644
--- a/onethingai.py
+++ b/onethingai.py
@@ -17,8 +17,9 @@ class OneThingAILoader:
                 "prompt": ("STRING", {"default": "", "multiline": True}),
                 "model": ("STRING", {"default": "gpt-image-1"}),
                 "num_images": ("INT", {"default": 1, "min": 1, "max": 10}),
-                "width": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
-                "height": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
+                "image_size": (["1536x1024", "1024x1024", "1024x1536", "自定义"], {"default": "1024x1024"}),
+                "custom_width": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
+                "custom_height": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
                 "retries": ("INT", {"default": 3, "min": 1, "max": 10}),
                 "timeout": ("INT", {"default": 120, "min": 5, "max": 180}),
             },
@@ -32,7 +33,7 @@ class OneThingAILoader:
     FUNCTION = "generate"
     CATEGORY = "OneThingAI/image generation"
 
-    def generate(self, api_key, prompt, model="gpt4o", num_images=1, width=1024, height=1024, retries=3, timeout=8, reference_image=None, reference_image_weight=0.5):
+    def generate(self, api_key, prompt, model="gpt-image-1", num_images=1, image_size="1024x1024", custom_width=1024, custom_height=1024, retries=3, timeout=8, reference_image=None, reference_image_weight=0.5):
         # API endpoint
         url = "https://api-model.onethingai.com/v1/images/generations"
         
@@ -54,12 +55,16 @@ class OneThingAILoader:
             "Content-Type": "application/json"
         }
         
+        # Determine image size based on selection
+        if image_size == "自定义":
+            image_size = f"{custom_width}x{custom_height}"
+
         # Request payload
         payload = {
             "model": model,
             "prompt": prompt,
             "n": num_images,
-            "size": f"{width}x{height}",
+            "size": size,
         }
 
         # Add reference image if provided