diff --git a/image_generation.py b/image_generation.py
index a55e3c1..e4e2a42 100644
--- a/image_generation.py
+++ b/image_generation.py
@@ -8,7 +8,7 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 class DeepSeekImageGeneration:
-    "Enhanced Image Generation Node for DeepSeek Janus Pro"
+    """Enhanced Image Generation Node for DeepSeek Janus Pro"""
     
     @classmethod
     def INPUT_TYPES(cls):
@@ -110,7 +110,7 @@ class DeepSeekImageGeneration:
             # Prepare inputs
             inputs = tokenizer(
                 conversations=conversation,
-                return_tensors='pt',
+                return_tensors="pt",
                 truncation=True,
                 max_length=max_length,
                 padding=True
@@ -122,18 +122,32 @@ class DeepSeekImageGeneration:
                 for k, v in inputs.items()
             }
 
-            # Generate images
-            with torch.no_grad():
-                outputs = model.generate(
-                    **batch_inputs,
-                    do_sample=True,
-                    temperature=temperature,
-                    guidance_scale=guidance_scale,
-                    num_beams=batch_size * 2,
-                    max_length=max_length + 20,
-                    pad_token_id=tokenizer.tokenizer.pad_token_id,
-                    eos_token_id=tokenizer.tokenizer.eos_token_id,
-                )
+            try:
+                # Generate images
+                with torch.no_grad():
+                    outputs = model.text_generate(  # Changed from generate to text_generate
+                        **batch_inputs,
+                        do_sample=True,
+                        temperature=temperature,
+                        guidance_scale=guidance_scale,
+                        num_beams=batch_size * 2,
+                        max_length=max_length + 20,
+                        pad_token_id=tokenizer.tokenizer.pad_token_id,
+                        eos_token_id=tokenizer.tokenizer.eos_token_id,
+                    )
+            except AttributeError:
+                # Fallback to regular generate if text_generate is not available
+                with torch.no_grad():
+                    outputs = model.generate(
+                        **batch_inputs,
+                        do_sample=True,
+                        temperature=temperature,
+                        guidance_scale=guidance_scale,
+                        num_beams=batch_size * 2,
+                        max_length=max_length + 20,
+                        pad_token_id=tokenizer.tokenizer.pad_token_id,
+                        eos_token_id=tokenizer.tokenizer.eos_token_id,
+                    )
 
             # Decode and process images
             images = model.decode_images(outputs)
@@ -171,4 +185,4 @@ NODE_CLASS_MAPPINGS = {
 # Node display name mappings
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DeepSeekImageGeneration": "DeepSeek Image Generation Pro"
-}
\ No newline at end of file
+}
diff --git a/image_understanding.py b/image_understanding.py
index ef8d0b2..16e1036 100644
--- a/image_understanding.py
+++ b/image_understanding.py
@@ -43,7 +43,7 @@ class DeepSeekImageUnderstanding:
         """
         try:
             # Log the image parameter
-            logger.info(f"Received image: {image}")
+            logger.info(f"Received image tensor with shape: {image.shape}")
 
             # 图像预处理
             if len(image.shape) == 4:  # BCHW format
@@ -72,12 +72,35 @@ class DeepSeekImageUnderstanding:
                 return_tensors='pt'
             ).to(model.device)
 
-            # 生成响应
-            outputs = model.generate(
-                **model_inputs,
-                max_new_tokens=512,
-                do_sample=False
-            )
+            try:
+                # 首先尝试使用 text_generate
+                with torch.no_grad():
+                    outputs = model.text_generate(
+                        **model_inputs,
+                        max_new_tokens=512,
+                        do_sample=False
+                    )
+            except AttributeError:
+                try:
+                    # 如果 text_generate 不可用，尝试使用 chat
+                    with torch.no_grad():
+                        outputs = model.chat(
+                            tokenizer,
+                            question,
+                            history=[],
+                            images=[pil_image],
+                            max_new_tokens=512,
+                            do_sample=False
+                        )
+                        return (outputs,)
+                except AttributeError:
+                    # 如果 chat 也不可用，使用标准 generate
+                    with torch.no_grad():
+                        outputs = model.generate(
+                            **model_inputs,
+                            max_new_tokens=512,
+                            do_sample=False
+                        )
 
             # 解码响应
             answer = tokenizer.decode(outputs[0], skip_special_tokens=True)