Stream recognition results with dual display: markdown text during recognition, render after completion. Args: input_image: Input PIL Image max_length: Maximum generation length
(input_image, max_length=2048)
| 74 | |
| 75 | # --- 2. Streaming generation function --- |
| 76 | def stream_recognize_image(input_image, max_length=2048): |
| 77 | """Stream recognition results with dual display: markdown text during recognition, render after completion. |
| 78 | |
| 79 | Args: |
| 80 | input_image: Input PIL Image |
| 81 | max_length: Maximum generation length |
| 82 | """ |
| 83 | if input_image is None: |
| 84 | yield 'Please upload an image first.', '**Please upload an image first.**' |
| 85 | return |
| 86 | |
| 87 | # Convert to PIL Image if needed |
| 88 | if not isinstance(input_image, Image.Image): |
| 89 | input_image = Image.fromarray(input_image).convert('RGB') |
| 90 | else: |
| 91 | input_image = input_image.convert('RGB') |
| 92 | |
| 93 | # Get token IDs |
| 94 | bos_token_id = model.tokenizer.bos_token_id |
| 95 | eos_token_id = model.tokenizer.eos_token_id |
| 96 | pad_token_id = model.tokenizer.pad_token_id |
| 97 | |
| 98 | # Encode image |
| 99 | encoder_hidden_states, cross_k, cross_v = model.encode_image(input_image) |
| 100 | |
| 101 | # Initialize generation |
| 102 | generated_ids = [bos_token_id] |
| 103 | |
| 104 | # Initialize empty past_key_values |
| 105 | batch_size = encoder_hidden_states.shape[0] |
| 106 | past_key_values = [] |
| 107 | for _ in range(model.num_decoder_layers): |
| 108 | empty_key = np.zeros( |
| 109 | (batch_size, model.num_heads, 0, model.head_dim), |
| 110 | dtype=np.float32) |
| 111 | empty_value = np.zeros( |
| 112 | (batch_size, model.num_heads, 0, model.head_dim), |
| 113 | dtype=np.float32) |
| 114 | past_key_values.append((empty_key, empty_value)) |
| 115 | |
| 116 | cleaned_text = '' |
| 117 | |
| 118 | # Generation loop with streaming |
| 119 | for step in range(max_length - 1): |
| 120 | current_token = generated_ids[-1] |
| 121 | past_length = step |
| 122 | |
| 123 | # Decode step |
| 124 | logits, past_key_values = model.decode_step( |
| 125 | current_token, |
| 126 | past_length, |
| 127 | cross_k, |
| 128 | cross_v, |
| 129 | past_key_values, |
| 130 | padding_idx=pad_token_id |
| 131 | ) |
| 132 | |
| 133 | # Get next token |
nothing calls this directly
no test coverage detected