fix zerogpu error

app.py

@@ -41,22 +41,46 @@ def _get_args():
                         action='store_true',
                         default=False,
                         help='Automatically launch the interface in a new tab on the default browser.')
-    # parser.add_argument('--server-port', type=int, default=8080, help='Demo server port.')
-    # parser.add_argument('--server-name', type=str, default='29.210.129.176', help='Demo server name.')
+    
 
     args = parser.parse_args()
     return args
 
 
 def _load_model_processor(args):
-    model = HunYuanVLForConditionalGeneration.from_pretrained(
-        args.checkpoint_path,
-        attn_implementation="eager",  # "flash_attention_2", #也可以是 flash_attention_2 或 sdpa，根据你的环境支持情况选择
-        torch_dtype=torch.bfloat16,
-        # device_map="auto",
-        device_map="cuda",
-        token=os.environ.get('HF_TOKEN')
-    )
+    # 优化：尝试使用 flash_attention_2 或 sdpa
+    try:
+        attn_impl = "flash_attention_2"
+        print(f"[INFO] 尝试使用 {attn_impl}")
+        model = HunYuanVLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            attn_implementation=attn_impl,
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
+            token=os.environ.get('HF_TOKEN')
+        )
+    except Exception as e:
+        print(f"[WARNING] flash_attention_2 不可用: {e}")
+        print(f"[INFO] 降级使用 sdpa")
+        try:
+            model = HunYuanVLForConditionalGeneration.from_pretrained(
+                args.checkpoint_path,
+                attn_implementation="sdpa",
+                torch_dtype=torch.bfloat16,
+                device_map="cuda",
+                token=os.environ.get('HF_TOKEN')
+            )
+        except Exception as e2:
+            print(f"[WARNING] sdpa 不可用: {e2}")
+            print(f"[INFO] 使用 eager (最慢)")
+            model = HunYuanVLForConditionalGeneration.from_pretrained(
+                args.checkpoint_path,
+                attn_implementation="eager",
+                torch_dtype=torch.bfloat16,
+                device_map="cuda",
+                token=os.environ.get('HF_TOKEN')
+            )
+    
     processor = AutoProcessor.from_pretrained(args.checkpoint_path, use_fast=False, trust_remote_code=True)
     return model, processor
 
@@ -88,16 +112,29 @@ def _gc():
 
 
 def _launch_demo(args, model, processor):
-    @spaces.GPU(duration=200)
+    # 关键修复：减少 duration，添加调试信息
+    @spaces.GPU(duration=60)
     def call_local_model(model, processor, messages):
-        print(messages)
+        import time
+        start_time = time.time()
+        print(f"[DEBUG] 开始推理，时间: {start_time}")
+        print(f"[DEBUG] Messages: {messages}")
+        
         messages = [messages]
         # 使用 processor 构造输入格式
         texts = [
             processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
             for msg in messages
         ]
+        
+        prep_time = time.time()
+        print(f"[DEBUG] 模板处理耗时: {prep_time - start_time:.2f}s")
+        
         image_inputs, video_inputs = process_vision_info(messages)
+        
+        vision_time = time.time()
+        print(f"[DEBUG] 视觉处理耗时: {vision_time - prep_time:.2f}s")
+        
         inputs = processor(
             text=texts,
             images=image_inputs,
@@ -107,25 +144,30 @@ def _launch_demo(args, model, processor):
         )
         inputs = inputs.to(model.device)
         
-
-        # gen_kwargs = {'max_new_tokens': 32768, 'streamer': streamer, **inputs}
-        # thread = Thread(target=model.generate, kwargs=gen_kwargs)
-        # thread.start()
-
-        # generated_text = ''
-        # for new_text in streamer:
-        #     generated_text += new_text
-        #     yield generated_text
-
-        # 模型推理
+        input_time = time.time()
+        print(f"[DEBUG] 输入处理耗时: {input_time - vision_time:.2f}s")
+        print(f"[DEBUG] Input shape: {inputs.input_ids.shape if 'input_ids' in inputs else 'N/A'}")
+        
+        # 关键修复1: 大幅减少 max_new_tokens
+        # 关键修复2: 添加 EOS token 和停止条件
+        # 关键修复3: 添加超时保护
         with torch.no_grad():
             generated_ids = model.generate(
                 **inputs,
-                max_new_tokens=1024*8,
+                max_new_tokens=512,  # 从 8192 降到 512，避免无限生成
                 repetition_penalty=1.03,
-                do_sample=False
+                do_sample=False,
+                # 关键：设置 EOS token，确保能正常停止
+                eos_token_id=processor.tokenizer.eos_token_id,
+                pad_token_id=processor.tokenizer.pad_token_id,
+                # 添加提前停止条件
+                use_cache=True,
             )
         
+        gen_time = time.time()
+        print(f"[DEBUG] 生成耗时: {gen_time - input_time:.2f}s")
+        print(f"[DEBUG] Generated shape: {generated_ids.shape}")
+        
         # 解码输出
         if "input_ids" in inputs:
             input_ids = inputs.input_ids
@@ -135,11 +177,18 @@ def _launch_demo(args, model, processor):
         generated_ids_trimmed = [
             out_ids[len(in_ids):] for in_ids, out_ids in zip(input_ids, generated_ids)
         ]
+        
+        print(f"[DEBUG] Trimmed tokens count: {[len(ids) for ids in generated_ids_trimmed]}")
 
         output_texts = processor.batch_decode(
             generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )
-
+        
+        decode_time = time.time()
+        print(f"[DEBUG] 解码耗时: {decode_time - gen_time:.2f}s")
+        print(f"[DEBUG] 总耗时: {decode_time - start_time:.2f}s")
+        print(f"[DEBUG] Output: {output_texts[0][:200]}...")  # 只打印前200字符
+        
         return output_texts