fix zerogpu error

app.py

@@ -48,40 +48,18 @@ def _get_args():
 
 
 def _load_model_processor(args):
-    # 优化：尝试使用 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')
-            )
-    
+    # ZeroGPU 环境：模型在 CPU 上加载，使用 eager 模式
+    # 在 @spaces.GPU 装饰器内会自动移到 GPU
+    print(f"[INFO] 加载模型（ZeroGPU 环境使用 eager 模式）")
+    model = HunYuanVLForConditionalGeneration.from_pretrained(
+        args.checkpoint_path,
+        attn_implementation="eager",  # ZeroGPU 必须用 eager，因为初始在 CPU
+        torch_dtype=torch.bfloat16,
+        device_map="auto",  # 改回 auto，让 ZeroGPU 自动管理
+        token=os.environ.get('HF_TOKEN')
+    )
     processor = AutoProcessor.from_pretrained(args.checkpoint_path, use_fast=False, trust_remote_code=True)
+    print(f"[INFO] 模型加载完成")
     return model, processor
 
 
@@ -112,28 +90,25 @@ def _gc():
 
 
 def _launch_demo(args, model, processor):
-    # 关键修复：减少 duration，添加调试信息
-    @spaces.GPU(duration=60)
+    # 关键：减少 duration 到 30 秒，如果超时说明有问题
+    @spaces.GPU(duration=30)
     def call_local_model(model, processor, messages):
         import time
         start_time = time.time()
-        print(f"[DEBUG] 开始推理，时间: {start_time}")
-        print(f"[DEBUG] Messages: {messages}")
+        print(f"[DEBUG] ========== 开始推理 ==========")
+        print(f"[DEBUG] 时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
         
         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")
+        print(f"[DEBUG] 模板构建完成，耗时: {time.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")
+        print(f"[DEBUG] 图像处理完成，耗时: {time.time() - start_time:.2f}s")
         
         inputs = processor(
             text=texts,
@@ -143,51 +118,55 @@ def _launch_demo(args, model, processor):
             return_tensors="pt",
         )
         inputs = inputs.to(model.device)
+        print(f"[DEBUG] 输入准备完成，耗时: {time.time() - start_time:.2f}s")
+        print(f"[DEBUG] Input IDs shape: {inputs.input_ids.shape}")
+        print(f"[DEBUG] Device: {model.device}")
         
-        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: 添加超时保护
+        # 关键优化：极限压缩参数
+        gen_start = time.time()
         with torch.no_grad():
             generated_ids = model.generate(
                 **inputs,
                 max_new_tokens=512,  # 从 8192 降到 512，避免无限生成
                 repetition_penalty=1.03,
                 do_sample=False,
-                # 关键：设置 EOS token，确保能正常停止
                 eos_token_id=processor.tokenizer.eos_token_id,
                 pad_token_id=processor.tokenizer.pad_token_id,
-                # 添加提前停止条件
                 use_cache=True,
+                # 关键：添加长度惩罚，鼓励短输出
+                length_penalty=0.8,
+                # 添加早停
+                early_stopping=True,
             )
         
-        gen_time = time.time()
-        print(f"[DEBUG] 生成耗时: {gen_time - input_time:.2f}s")
-        print(f"[DEBUG] Generated shape: {generated_ids.shape}")
+        gen_time = time.time() - gen_start
+        print(f"[DEBUG] ========== 生成完成 ==========")
+        print(f"[DEBUG] 生成耗时: {gen_time:.2f}s")
+        print(f"[DEBUG] Output shape: {generated_ids.shape}")
         
         # 解码输出
         if "input_ids" in inputs:
             input_ids = inputs.input_ids
         else:
-            input_ids = inputs.inputs  # fallback
+            input_ids = inputs.inputs
 
         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]}")
+        actual_tokens = len(generated_ids_trimmed[0])
+        print(f"[DEBUG] 实际生成 token 数: {actual_tokens}")
+        print(f"[DEBUG] 每 token 耗时: {gen_time/actual_tokens if actual_tokens > 0 else 0:.3f}s")
 
         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字符
+        total_time = time.time() - start_time
+        print(f"[DEBUG] ========== 全部完成 ==========")
+        print(f"[DEBUG] 总耗时: {total_time:.2f}s")
+        print(f"[DEBUG] 输出长度: {len(output_texts[0])} 字符")
+        print(f"[DEBUG] 输出预览: {output_texts[0][:100]}...")
         
         return output_texts