File size: 39,753 Bytes
5d5f953
 
 
 
7352136
5d5f953
 
 
 
 
 
 
 
 
 
 
7352136
bdbf47f
 
 
 
 
 
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1efad72
5d5f953
 
 
 
 
 
acfce9f
 
 
c2b0812
 
acfce9f
 
 
 
 
 
bdbf47f
 
 
 
 
 
 
 
 
 
5d5f953
c2b0812
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e220a3a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d5f953
 
bdbf47f
 
 
d87f42b
c2b0812
 
d87f42b
1efad72
bdbf47f
1efad72
bdbf47f
 
 
 
 
 
 
acfce9f
bdbf47f
 
c2b0812
 
 
 
 
 
 
d87f42b
 
 
 
c2b0812
d87f42b
 
 
c2b0812
 
 
 
 
 
 
 
 
 
 
1efad72
5d5f953
acfce9f
5d5f953
 
 
 
 
acfce9f
1efad72
5d5f953
acfce9f
1efad72
bdbf47f
 
 
 
 
 
 
 
 
 
 
 
 
5d5f953
 
 
 
 
 
 
bdbf47f
 
d87f42b
bdbf47f
d87f42b
bdbf47f
 
acfce9f
d87f42b
e7257d2
5d5f953
d87f42b
acfce9f
e7257d2
 
 
 
fe3cde8
e7257d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d5f953
bdbf47f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e7257d2
bdbf47f
e7257d2
 
f72db18
c8398d5
e7257d2
bdbf47f
e7257d2
 
bdbf47f
 
e7257d2
 
bdbf47f
5d5f953
acfce9f
 
 
 
1efad72
5d5f953
 
 
 
acfce9f
5d5f953
 
 
 
1efad72
acfce9f
 
 
5d5f953
 
 
 
1efad72
e220a3a
acfce9f
 
 
 
 
e220a3a
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
95825d6
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d87f42b
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0b2eee
5d5f953
 
 
 
 
 
 
 
cfbd98d
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2852b4
 
 
 
9355006
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a88fed3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
95825d6
 
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2852b4
5d5f953
 
 
 
 
 
 
 
 
a2852b4
5d5f953
 
 
 
 
 
 
 
 
a2852b4
5d5f953
 
 
 
 
 
 
 
 
a2852b4
5d5f953
 
 
 
 
 
 
 
 
9355006
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9355006
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fe3cde8
 
5d5f953
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
import os
import torch
import numpy as np
from PIL import Image
import spaces
from transformers import AutoProcessor
from qwen_vl_utils import process_vision_info  # 请确保该模块在你的环境可用
from transformers import HunYuanVLForConditionalGeneration
import gradio as gr
from argparse import ArgumentParser
import copy
import requests
from io import BytesIO
import tempfile
import hashlib
import gc

# 关键优化:设置环境变量加速 transformers
os.environ['TOKENIZERS_PARALLELISM'] = 'false'  # 避免tokenizer警告
os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = '1'
# 禁用 PyTorch 的 JIT 融合优化(在某些情况下会导致首次运行极慢)
# torch._C._jit_set_profiling_executor(False)
# torch._C._jit_set_profiling_mode(False)





def _get_args():
    parser = ArgumentParser()

    parser.add_argument('-c',
                        '--checkpoint-path',
                        type=str,
                        default='tencent/HunyuanOCR',
                        help='Checkpoint name or path, default to %(default)r')
    parser.add_argument('--cpu-only', action='store_true', help='Run demo with CPU only')

    parser.add_argument('--flash-attn2',
                        action='store_true',
                        default=False,
                        help='Enable flash_attention_2 when loading the model.')
    parser.add_argument('--share',
                        action='store_true',
                        default=False,
                        help='Create a publicly shareable link for the interface.')
    parser.add_argument('--inbrowser',
                        action='store_true',
                        default=False,
                        help='Automatically launch the interface in a new tab on the default browser.')
    

    args = parser.parse_args()
    return args


def _load_model_processor(args):
    # ZeroGPU 环境:模型在 CPU 上加载,使用 eager 模式
    # 在 @spaces.GPU 装饰器内会自动移到 GPU
    print(f"[INFO] 加载模型(ZeroGPU 环境使用 eager 模式)")
    print(f"[INFO] 加载时 CUDA available: {torch.cuda.is_available()}")
    
    model = HunYuanVLForConditionalGeneration.from_pretrained(
        args.checkpoint_path,
        attn_implementation="eager",  # ZeroGPU 必须用 eager,因为初始在 CPU
        torch_dtype=torch.bfloat16,
        device_map="auto",  # 改回 auto,让 ZeroGPU 自动管理
    )
    
    # 关键:禁用梯度检查点(如果启用会导致极慢)
    if hasattr(model, 'gradient_checkpointing_disable'):
        model.gradient_checkpointing_disable()
        print(f"[INFO] 梯度检查点已禁用")
    
    # 设置为评估模式
    model.eval()
    print(f"[INFO] 模型设置为评估模式")
    
    processor = AutoProcessor.from_pretrained(args.checkpoint_path, use_fast=False, trust_remote_code=True)
    
    print(f"[INFO] 模型加载完成,当前设备: {next(model.parameters()).device}")
    return model, processor


def _parse_text(text):
    """解析文本,处理特殊格式"""
    # if text is None:
    #     return text
    text = text.replace("<trans>", "").replace("</trans>", "")
    return text


def _remove_image_special(text):
    """移除图像特殊标记"""
    # if text is None:
    #     return text
    # # 移除可能的图像特殊标记
    # import re
    # text = re.sub(r'<image>|</image>|<img>|</img>', '', text)
    # return text
    return text


def _gc():
    """垃圾回收"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def clean_repeated_substrings(text):
    """Clean repeated substrings in text"""
    n = len(text)
    if n < 2000:
        return text
    for length in range(2, n // 10 + 1):
        candidate = text[-length:] 
        count = 0
        i = n - length
        
        while i >= 0 and text[i:i + length] == candidate:
            count += 1
            i -= length

        if count >= 10:
            return text[:n - length * (count - 1)]  

    return text


def _launch_demo(args, model, processor):
    # 全局变量用于跟踪是否是首次调用
    first_call = [True]
    
    # 关键修复:移除 model 和 processor 参数,使用闭包访问
    # 增加 duration 到 120 秒,避免高峰期超时
    @spaces.GPU(duration=120)
    def call_local_model(messages):
        import time
        import sys
        start_time = time.time()
        
        if first_call[0]:
            print(f"[INFO] ========== 这是首次推理调用 ==========")
            first_call[0] = False
        else:
            print(f"[INFO] ========== 这是第 N 次推理调用 ==========")
        
        print(f"[DEBUG] ========== 开始推理 ==========")
        print(f"[DEBUG] Python version: {sys.version}")
        print(f"[DEBUG] PyTorch version: {torch.__version__}")
        print(f"[DEBUG] CUDA available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"[DEBUG] CUDA device count: {torch.cuda.device_count()}")
            print(f"[DEBUG] Current CUDA device: {torch.cuda.current_device()}")
            print(f"[DEBUG] Device name: {torch.cuda.get_device_name(0)}")
            print(f"[DEBUG] GPU Memory allocated: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
            print(f"[DEBUG] GPU Memory reserved: {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")
        
        # 关键:检查并确保模型在 GPU 上
        model_device = next(model.parameters()).device
        print(f"[DEBUG] Model device: {model_device}")
        print(f"[DEBUG] Model dtype: {next(model.parameters()).dtype}")
        
        if str(model_device) == 'cpu':
            print(f"[ERROR] 模型在 CPU 上!尝试移动到 GPU...")
            if torch.cuda.is_available():
                move_start = time.time()
                model.cuda()
                move_time = time.time() - move_start
                print(f"[DEBUG] Model device after cuda(): {next(model.parameters()).device}")
                print(f"[DEBUG] 模型移动到 GPU 耗时: {move_time:.2f}s")
            else:
                print(f"[CRITICAL] CUDA 不可用!将在 CPU 上运行,速度会很慢!")
                print(f"[CRITICAL] 这可能是因为 ZeroGPU 资源紧张或超时")
        else:
            print(f"[INFO] 模型已在 GPU 上: {model_device}")
        
        messages = [messages]
        
        # 使用 processor 构造输入格式
        texts = [
            processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
            for msg in messages
        ]
        print(f"[DEBUG] 模板构建完成,耗时: {time.time() - start_time:.2f}s")
        
        image_inputs, video_inputs = process_vision_info(messages)
        print(f"[DEBUG] 图像处理完成,耗时: {time.time() - start_time:.2f}s")
        
        # 检查图像输入大小
        if image_inputs:
            for idx, img in enumerate(image_inputs):
                if hasattr(img, 'size'):
                    print(f"[DEBUG] Image {idx} size: {img.size}")
                elif isinstance(img, np.ndarray):
                    print(f"[DEBUG] Image {idx} shape: {img.shape}")
        
        print(f"[DEBUG] 开始 processor 编码输入...")
        processor_start = time.time()
        
        print(f"[DEBUG] 开始 processor 编码输入...")
        processor_start = time.time()
        inputs = processor(
            text=texts,
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        print(f"[DEBUG] Processor 编码完成,耗时: {time.time() - processor_start:.2f}s")
        
        # 确保输入在 GPU 上
        to_device_start = time.time()
        inputs = inputs.to('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"[DEBUG] 输入移到设备耗时: {time.time() - to_device_start:.2f}s")
        print(f"[DEBUG] 输入准备完成,总耗时: {time.time() - start_time:.2f}s")
        print(f"[DEBUG] Input IDs shape: {inputs.input_ids.shape}")
        print(f"[DEBUG] Input device: {inputs.input_ids.device}")
        print(f"[DEBUG] Input sequence length: {inputs.input_ids.shape[1]}")
        
        # 生成
        gen_start = time.time()
        print(f"[DEBUG] ========== 开始生成 tokens ==========")
        
        # 关键优化:根据任务类型动态调整 max_new_tokens
        # OCR 任务通常不需要 8192 tokens,这会导致不必要的等待
        max_new_tokens = 2048  # 从 8192 降到 2048,大幅提速
        print(f"[DEBUG] max_new_tokens: {max_new_tokens}")
        
        # 添加进度回调
        token_count = [0]
        last_time = [gen_start]
        
        def progress_callback(input_ids, scores, **kwargs):
            token_count[0] += 1
            current_time = time.time()
            if token_count[0] % 10 == 0 or (current_time - last_time[0]) > 2.0:
                elapsed = current_time - gen_start
                tokens_per_sec = token_count[0] / elapsed if elapsed > 0 else 0
                print(f"[DEBUG] 已生成 {token_count[0]} tokens, 速度: {tokens_per_sec:.2f} tokens/s, 耗时: {elapsed:.2f}s")
                last_time[0] = current_time
            return False
        
        with torch.no_grad():
            print(f"[DEBUG] 进入 torch.no_grad() 上下文,耗时: {time.time() - start_time:.2f}s")
            
            # 先做一次简单的前向传播测试
            print(f"[DEBUG] 测试前向传播...")
            forward_test_start = time.time()
            try:
                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                    test_outputs = model(**inputs, use_cache=False)
                print(f"[DEBUG] 前向传播测试成功,耗时: {time.time() - forward_test_start:.2f}s")
            except Exception as e:
                print(f"[WARNING] 前向传播测试失败: {e}")
            
            print(f"[DEBUG] 开始调用 model.generate()... (当前耗时: {time.time() - start_time:.2f}s)")
            generate_call_start = time.time()
            
            try:
                # 关键:添加更激进的生成参数,强制早停
                generated_ids = model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                )
                print(f"[DEBUG] model.generate() 返回,耗时: {time.time() - generate_call_start:.2f}s")
            except Exception as e:
                print(f"[ERROR] 生成失败: {e}")
                import traceback
                traceback.print_exc()
                raise
        
        print(f"[DEBUG] 退出 torch.no_grad() 上下文")
        
        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

        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(input_ids, generated_ids)
        ]
        
        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
        )
        
        
        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]}...")
        output_texts[0] = clean_repeated_substrings(output_texts[0])
        return output_texts
    

    def create_predict_fn():

        def predict(_chatbot, task_history):
            nonlocal model, processor
            chat_query = _chatbot[-1][0]
            query = task_history[-1][0]
            if len(chat_query) == 0:
                _chatbot.pop()
                task_history.pop()
                return _chatbot
            print('User: ', query)
            history_cp = copy.deepcopy(task_history)
            full_response = ''
            messages = []
            messages.append({'role': 'system', 'content': ""})
            content = []
            for q, a in history_cp:
                if isinstance(q, (tuple, list)):
                    # 判断是URL还是本地路径
                    img_path = q[0]
                    if img_path.startswith(('http://', 'https://')):
                        content.append({'type': 'image', 'image': img_path})
                    else:
                        content.append({'type': 'image', 'image': f'{os.path.abspath(img_path)}'})
                else:
                    content.append({'type': 'text', 'text': q})
                    messages.append({'role': 'user', 'content': content})
                    messages.append({'role': 'assistant', 'content': [{'type': 'text', 'text': a}]})
                    content = []
            messages.pop()
            
            # 调用模型获取响应(已修改:不再传递 model 和 processor)
            response_list = call_local_model(messages)
            response = response_list[0] if response_list else ""
            
            _chatbot[-1] = (_parse_text(chat_query), _remove_image_special(_parse_text(response)))
            full_response = _parse_text(response)

            task_history[-1] = (query, full_response)
            print('HunyuanOCR: ' + _parse_text(full_response))
            yield _chatbot

        return predict
    
    def create_regenerate_fn():

        def regenerate(_chatbot, task_history):
            nonlocal model, processor
            if not task_history:
                return _chatbot
            item = task_history[-1]
            if item[1] is None:
                return _chatbot
            task_history[-1] = (item[0], None)
            chatbot_item = _chatbot.pop(-1)
            if chatbot_item[0] is None:
                _chatbot[-1] = (_chatbot[-1][0], None)
            else:
                _chatbot.append((chatbot_item[0], None))
            # 使用外层的predict函数
            _chatbot_gen = predict(_chatbot, task_history)
            for _chatbot in _chatbot_gen:
                yield _chatbot

        return regenerate

    predict = create_predict_fn()
    regenerate = create_regenerate_fn()

    def add_text(history, task_history, text):
        task_text = text
        history = history if history is not None else []
        task_history = task_history if task_history is not None else []
        history = history + [(_parse_text(text), None)]
        task_history = task_history + [(task_text, None)]
        return history, task_history, ''

    def add_file(history, task_history, file):
        history = history if history is not None else []
        task_history = task_history if task_history is not None else []
        history = history + [((file.name,), None)]
        task_history = task_history + [((file.name,), None)]
        return history, task_history
    
    def download_url_image(url):
        """下载 URL 图片到本地临时文件"""
        try:
            # 使用 URL 的哈希值作为文件名,避免重复下载
            url_hash = hashlib.md5(url.encode()).hexdigest()
            temp_dir = tempfile.gettempdir()
            temp_path = os.path.join(temp_dir, f"hyocr_demo_{url_hash}.png")
            
            # 如果文件已存在,直接返回
            if os.path.exists(temp_path):
                return temp_path
            
            # 下载图片
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            with open(temp_path, 'wb') as f:
                f.write(response.content)
            return temp_path
        except Exception as e:
            print(f"下载图片失败: {url}, 错误: {e}")
            return url  # 失败时返回原 URL

    def reset_user_input():
        return gr.update(value='')

    def reset_state(_chatbot, task_history):
        task_history.clear()
        _chatbot.clear()
        _gc()
        return []

    # 示例图片路径配置 - 请替换为实际图片路径
    EXAMPLE_IMAGES = {
        "spotting": "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/23cc43af9376b948f3febaf4ce854a8a.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523817%3B1794627877&q-key-time=1763523817%3B1794627877&q-header-list=host&q-url-param-list=&q-signature=8ebd6a9d3ed7eba73bb783c337349db9c29972e2",   # TODO: 替换为场景文字示例图片路径
        "parsing": "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/c4997ebd1be9f7c3e002fabba8b46cb7.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=d2cd12be4c7902821c8c82203e4642624046911a",
        "ie": "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/7c67c0f78e4423d51644a325da1f8e85.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=803648f3253706f654faf1423869fd9e00e7056e",
        "vqa": "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/fea0865d1c70c53aaa2ab91cd0e787f5.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=a92b94e298a11aea130d730d3b16ee761acc3f4c",
        "translation": "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/d1af99d35e9db9e820ebebb5bc68993a.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763967603%3B1795071663&q-key-time=1763967603%3B1795071663&q-header-list=host&q-url-param-list=&q-signature=a57080c0b3d4c76ea74b88c6291f9004241c9d49",
        # "spotting": "examples/spotting.jpg",  
        # "parsing": "examples/parsing.jpg", 
        # "ie": "examples/ie.jpg",   
        # "vqa": "examples/vqa.jpg",            
        # "translation": "examples/translation.jpg"  
    }

    with gr.Blocks(css="""
        body {
            background: #f5f7fa;
        }
        .gradio-container {
            max-width: 100% !important;
            padding: 0 40px !important;
        }
        .header-section {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            padding: 30px 0;
            margin: -20px -40px 30px -40px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        }
        .header-content {
            max-width: 1600px;
            margin: 0 auto;
            padding: 0 40px;
            display: flex;
            align-items: center;
            gap: 20px;
        }
        .header-logo {
            height: 60px;
        }
        .header-text h1 {
            color: white;
            font-size: 32px;
            font-weight: bold;
            margin: 0 0 5px 0;
        }
        .header-text p {
            color: rgba(255,255,255,0.9);
            margin: 0;
            font-size: 14px;
        }
        .main-container {
            max-width: 1800px;
            margin: 0 auto;
        }
        .chatbot {
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08) !important;
            border-radius: 12px !important;
            border: 1px solid #e5e7eb !important;
            background: white !important;
        }
        .input-panel {
            background: white;
            padding: 20px;
            border-radius: 12px;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
            border: 1px solid #e5e7eb;
        }
        .input-box textarea {
            border: 2px solid #e5e7eb !important;
            border-radius: 8px !important;
            font-size: 14px !important;
        }
        .input-box textarea:focus {
            border-color: #667eea !important;
        }
        .btn-primary {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
            border: none !important;
            color: white !important;
            font-weight: 500 !important;
            padding: 10px 24px !important;
            font-size: 14px !important;
        }
        .btn-primary:hover {
            transform: translateY(-1px) !important;
            box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important;
        }
        .btn-secondary {
            background: white !important;
            border: 2px solid #667eea !important;
            color: #667eea !important;
            padding: 8px 20px !important;
            font-size: 14px !important;
        }
        .btn-secondary:hover {
            background: #f0f4ff !important;
        }
        .example-grid {
            display: grid;
            grid-template-columns: repeat(4, 1fr);
            gap: 20px;
            margin-top: 30px;
        }
        .example-card {
            background: white;
            border-radius: 12px;
            overflow: hidden;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
            border: 1px solid #e5e7eb;
            transition: all 0.3s ease;
        }
        .example-card:hover {
            transform: translateY(-4px);
            box-shadow: 0 8px 20px rgba(102, 126, 234, 0.15);
            border-color: #667eea;
        }
        .example-image-wrapper {
            width: 100%;
            height: 180px;
            overflow: hidden;
            background: #f5f7fa;
        }
        .example-image-wrapper img {
            width: 100%;
            height: 100%;
            object-fit: cover;
        }
        .example-btn {
            width: 100% !important;
            white-space: pre-wrap !important;
            text-align: left !important;
            padding: 16px !important;
            background: white !important;
            border: none !important;
            border-top: 1px solid #e5e7eb !important;
            color: #1f2937 !important;
            font-size: 14px !important;
            line-height: 1.6 !important;
            transition: all 0.3s ease !important;
            font-weight: 500 !important;
        }
        .example-btn:hover {
            background: #f9fafb !important;
            color: #667eea !important;
        }
        .feature-section {
            background: white;
            padding: 24px;
            border-radius: 12px;
            margin-top: 30px;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
            border: 1px solid #e5e7eb;
        }
        .section-title {
            font-size: 18px;
            font-weight: 600;
            color: #1f2937;
            margin-bottom: 20px;
            padding-bottom: 12px;
            border-bottom: 2px solid #e5e7eb;
        }
    """) as demo:
        # 顶部导航栏
        gr.HTML("""
        <div class="header-section">
            <div class="header-content">
                <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/6ef6928b21b323b2b00115f86a779d8f.png?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763450355%3B1794554415&q-key-time=1763450355%3B1794554415&q-header-list=host&q-url-param-list=&q-signature=41328696dc34571324aa18c791c1196192e729c6" class="header-logo"/>
                <div class="header-text">
                    <h1>HunyuanOCR</h1>
                    <p>Powered by Tencent Hunyuan Team</p>
                </div>
            </div>
            <div style="max-width: 1600px; margin: 0 auto; padding: 0 40px;">
                <div style="background: linear-gradient(135deg, #fff3cd 0%, #ffeaa7 100%); 
                            border-left: 4px solid #ffc107; 
                            padding: 16px 20px; 
                            margin: 20px 0 0 0; 
                            border-radius: 8px; 
                            box-shadow: 0 2px 8px rgba(255, 193, 7, 0.15);">
                    <div style="display: flex; align-items: flex-start; gap: 12px;">
                        <span style="font-size: 24px; flex-shrink: 0; line-height: 1;">⚠️</span>
                        <div style="flex: 1;">
                            <p style="margin: 0 0 8px 0; 
                                       font-size: 14px; 
                                       font-weight: 600; 
                                       color: #856404; 
                                       line-height: 1.65;">
                                <strong>🔔 重要提示:</strong>请注意,当前模型的 Transformers 实现精度尚未完全对齐(团队正在修复)。此 Space 旨在供您快速体验模型,若需获得完整精度与最佳性能,我们推荐使用 vLLM 进行部署。
                            </p>
                            <p style="margin: 0; 
                                       font-size: 13px; 
                                       font-weight: 500; 
                                       color: #856404; 
                                       line-height: 1.6;
                                       opacity: 0.92;">
                                <strong>🔔 Important Notice:</strong> Please note that the current Transformers implementation of the model has not yet achieved full precision alignment (the team is working on a fix). This Space is intended for quick model experimentation. For full precision and optimal performance, we recommend deploying via vLLM.
                            </p>
                        </div>
                    </div>
                </div>
            </div>
        </div>
        """)
        
        with gr.Column(elem_classes=["main-container"]):
            # 对话区域 - 全宽
            chatbot = gr.Chatbot(
                label='💬 对话窗口', 
                height=600,
                bubble_full_width=False,
                layout="bubble",
                show_copy_button=True,
                avatar_images=(None, "https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/6ef6928b21b323b2b00115f86a779d8f.png?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763450355%3B1794554415&q-key-time=1763450355%3B1794554415&q-header-list=host&q-url-param-list=&q-signature=41328696dc34571324aa18c791c1196192e729c6"),
                elem_classes=["chatbot"]
            )
            
            # 输入控制面板 - 全宽
            with gr.Group(elem_classes=["input-panel"]):
                query = gr.Textbox(
                    lines=2, 
                    label='💭 输入您的问题', 
                    placeholder='请先上传图片,然后输入问题。例如:检测并识别图片中的文字,将文本坐标格式化输出。',
                    elem_classes=["input-box"],
                    show_label=False
                )
                
                with gr.Row():
                    addfile_btn = gr.UploadButton('📁 上传图片', file_types=['image'], elem_classes=["btn-secondary"])
                    submit_btn = gr.Button('🚀 发送消息', variant="primary", elem_classes=["btn-primary"], scale=3)
                    regen_btn = gr.Button('🔄 重新生成', elem_classes=["btn-secondary"])
                    empty_bin = gr.Button('🗑️ 清空对话', elem_classes=["btn-secondary"])
            
            # 示例区域 - 5列网格布局
            gr.HTML('<div class="section-title">📚 快速体验示例 - 点击下方卡片快速加载</div>')
            
            with gr.Row():
                # 示例1:spotting
                with gr.Column(scale=1):
                    with gr.Group(elem_classes=["example-card"]):
                        gr.HTML("""
                        <div class="example-image-wrapper">
                            <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/23cc43af9376b948f3febaf4ce854a8a.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523817%3B1794627877&q-key-time=1763523817%3B1794627877&q-header-list=host&q-url-param-list=&q-signature=8ebd6a9d3ed7eba73bb783c337349db9c29972e2" alt="文字检测识别"/>
                        </div>
                        """)
                        example_1_btn = gr.Button("🔍 文字检测和识别", elem_classes=["example-btn"])
                
                # 示例2:parsing
                with gr.Column(scale=1):
                    with gr.Group(elem_classes=["example-card"]):
                        gr.HTML("""
                        <div class="example-image-wrapper">
                            <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/c4997ebd1be9f7c3e002fabba8b46cb7.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=d2cd12be4c7902821c8c82203e4642624046911a" alt="文档解析"/>
                        </div>
                        """)
                        example_2_btn = gr.Button("📋 文档解析", elem_classes=["example-btn"])
                
                # 示例3:ie
                with gr.Column(scale=1):
                    with gr.Group(elem_classes=["example-card"]):
                        gr.HTML("""
                        <div class="example-image-wrapper">
                            <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/7c67c0f78e4423d51644a325da1f8e85.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=803648f3253706f654faf1423869fd9e00e7056e" alt="信息抽取"/>
                        </div>
                        """)
                        example_3_btn = gr.Button("🎯 信息抽取", elem_classes=["example-btn"])
                
                # 示例4:VQA
                with gr.Column(scale=1):
                    with gr.Group(elem_classes=["example-card"]):
                        gr.HTML("""
                        <div class="example-image-wrapper">
                            <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/fea0865d1c70c53aaa2ab91cd0e787f5.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763523818%3B1794627878&q-key-time=1763523818%3B1794627878&q-header-list=host&q-url-param-list=&q-signature=a92b94e298a11aea130d730d3b16ee761acc3f4c" alt="视觉问答"/>
                        </div>
                        """)
                        example_4_btn = gr.Button("💬 视觉问答", elem_classes=["example-btn"])
                
                # 示例5:translation
                with gr.Column(scale=1):
                    with gr.Group(elem_classes=["example-card"]):
                        gr.HTML("""
                        <div class="example-image-wrapper">
                            <img src="https://hunyuan-multimodal-1258344703.cos.ap-guangzhou.myqcloud.com/hunyuan_multimodal/mllm_data/d1af99d35e9db9e820ebebb5bc68993a.jpg?q-sign-algorithm=sha1&q-ak=AKIDbLEFMUYZgyERZnygUQLC7xkQ1hTAzulX&q-sign-time=1763967603%3B1795071663&q-key-time=1763967603%3B1795071663&q-header-list=host&q-url-param-list=&q-signature=a57080c0b3d4c76ea74b88c6291f9004241c9d49" alt="图片翻译"/>
                        </div>
                        """)
                        example_5_btn = gr.Button("🌐 图片翻译", elem_classes=["example-btn"])
        
        task_history = gr.State([])
        
        
        # 示例1:文档识别
        def load_example_1(history, task_hist):
            prompt = "检测并识别图片中的文字,将文本坐标格式化输出。"
            image_url = EXAMPLE_IMAGES["spotting"]
            # 下载 URL 图片到本地
            image_path = download_url_image(image_url)
            # 清空对话历史
            history = []
            task_hist = []
            history = history + [((image_path,), None)]
            task_hist = task_hist + [((image_path,), None)]
            return history, task_hist, prompt
        
        
        
        # 示例2:场景文字
        def load_example_2(history, task_hist):
            prompt = "提取文档图片中正文的所有信息用markdown 格式表示,其中页眉、页脚部分忽略,表格用html 格式表达,文档中公式用latex 格式表示,按照阅读顺序组织进行解析。"
            image_url = EXAMPLE_IMAGES["parsing"]
            # 下载 URL 图片到本地
            image_path = download_url_image(image_url)
            # 清空对话历史
            history = []
            task_hist = []
            history = history + [((image_path,), None)]
            task_hist = task_hist + [((image_path,), None)]
            return history, task_hist, prompt
        
        
        
        # 示例3:表格提取
        def load_example_3(history, task_hist):
            prompt = "提取图片中的:['单价', '上车时间','发票号码', '省前缀', '总金额', '发票代码', '下车时间', '里程数'] 的字段内容,并且按照JSON格式返回。"
            image_url = EXAMPLE_IMAGES["ie"]
            # 下载 URL 图片到本地
            image_path = download_url_image(image_url)
            # 清空对话历史
            history = []
            task_hist = []
            history = history + [((image_path,), None)]
            task_hist = task_hist + [((image_path,), None)]
            return history, task_hist, prompt
        
        # 示例4:手写体
        def load_example_4(history, task_hist):
            prompt = "What is the highest life expectancy at birth of male?"
            image_url = EXAMPLE_IMAGES["vqa"]
            # 下载 URL 图片到本地
            image_path = download_url_image(image_url)
            # 清空对话历史
            history = []
            task_hist = []
            history = history + [((image_path,), None)]
            task_hist = task_hist + [((image_path,), None)]
            return history, task_hist, prompt
        
        # 示例5:翻译
        def load_example_5(history, task_hist):
            prompt = "将图中文字翻译为中文。"
            image_url = EXAMPLE_IMAGES["translation"]
            # 下载 URL 图片到本地
            image_path = download_url_image(image_url)
            # 清空对话历史
            history = []
            task_hist = []
            history = history + [((image_path,), None)]
            task_hist = task_hist + [((image_path,), None)]
            return history, task_hist, prompt
        
        # 绑定事件
        example_1_btn.click(load_example_1, [chatbot, task_history], [chatbot, task_history, query])
        example_2_btn.click(load_example_2, [chatbot, task_history], [chatbot, task_history, query])
        example_3_btn.click(load_example_3, [chatbot, task_history], [chatbot, task_history, query])
        example_4_btn.click(load_example_4, [chatbot, task_history], [chatbot, task_history, query])
        example_5_btn.click(load_example_5, [chatbot, task_history], [chatbot, task_history, query])
        
        submit_btn.click(add_text, [chatbot, task_history, query],
                         [chatbot, task_history]).then(predict, [chatbot, task_history], [chatbot], show_progress=True)
        submit_btn.click(reset_user_input, [], [query])
        empty_bin.click(reset_state, [chatbot, task_history], [chatbot], show_progress=True)
        regen_btn.click(regenerate, [chatbot, task_history], [chatbot], show_progress=True)
        addfile_btn.upload(add_file, [chatbot, task_history, addfile_btn], [chatbot, task_history], show_progress=True)

        # 功能说明区域
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML("""
                <div class="feature-section">
                    <div class="section-title">✨ 核心功能</div>
                    <ul style="line-height: 2; color: #4b5563; font-size: 14px; margin: 0; padding-left: 20px;">
                        <li><strong>🎯 高精度文字检测识别</strong> - 支持多场景文字检测与识别</li>
                        <li><strong>📐 智能文档解析</strong> - 自动识别文档结构,支持多粒度文档解析</li>
                        <li><strong>📋 信息提取</strong> - 支持30+高频卡证票据识别和结构化输出</li>
                        <li><strong>✏️ 视觉问答</strong> - 支持以文本为中心的开放式问答</li>
                        <li><strong>🌍 跨语言翻译</strong> - 支持中英互译及14+语种译为中英文</li>
                    </ul>
                </div>
                """)
            
            with gr.Column(scale=1):
                gr.HTML("""
                <div class="feature-section">
                    <div class="section-title">💡 使用建议</div>
                    <ul style="line-height: 2; color: #4b5563; font-size: 14px; margin: 0; padding-left: 20px;">
                        <li><strong>推理框架</strong> - 正式生产推荐使用VLLM,以获取更好的推理性能和精度</li>
                        <li><strong>拍摄角度</strong> - 确保图片清晰,光线充足,分辨率适中,避免严重倾斜、遮挡或反光,正面拍摄效果最佳</li>
                        <li><strong>文件大小</strong> - 建议单张图片不超过 10MB,支持 JPG/PNG 格式</li>
                        <li><strong>使用场景</strong> - 适用于文字检测识别、文档数字化、票据识别、信息提取、文字图片翻译等</li>
                        <li><strong>合规使用</strong> - 仅供学习研究,请遵守法律法规,尊重隐私权</li>
                    </ul>
                </div>
                """)
        
        # 底部版权信息
        gr.HTML("""
        <div style="text-align: center; color: #9ca3af; font-size: 13px; margin-top: 40px; padding: 20px; border-top: 1px solid #e5e7eb;">
            <p style="margin: 0;">© 2025 Tencent Hunyuan Team. All rights reserved.</p>
            <p style="margin: 5px 0 0 0;">本系统基于 HunyuanOCR 构建 | 仅供学习研究使用</p>
        </div>
        """)

    demo.queue().launch(
        share=args.share,
        inbrowser=args.inbrowser,
        # server_port=args.server_port,
        # server_name=args.server_name,
    )


def main():
    args = _get_args()
    model, processor = _load_model_processor(args)
    _launch_demo(args, model, processor)


if __name__ == '__main__':
    main()