asterocr 1.0.0__py3-none-any.whl

asterocr/kaggle_bridge/kaggle_worker.py

@@ -0,0 +1,1911 @@
+"""
+AsterOCR Kaggle Cloud Worker (Self-Contained)
+Runs on NVIDIA Tesla P100 (16GB VRAM) in Kaggle Cloud.
+Streams chunks from local machine via Cloudflare Tunnel, trains with FP16 Mixed Precision,
+streams real-time telemetry back to local Web GUI, and syncs checkpoints.
+"""
+
+import os
+import io
+import sys
+import time
+import json
+import struct
+import base64
+import gc
+import queue
+import threading
+import concurrent.futures
+import requests
+from PIL import Image
+import numpy as np
+import math
+
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.nn import functional as F
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import base64
+from io import BytesIO
+from PIL import Image
+
+def tensor_to_b64(tensor):
+    try:
+        arr = (tensor.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
+        if arr.shape[2] == 1:
+            arr = np.squeeze(arr, axis=2)
+        img = Image.fromarray(arr)
+        buffered = BytesIO()
+        img.save(buffered, format="JPEG", quality=60)
+        return "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode('utf-8')
+    except Exception:
+        return ""
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+# =========================================================================
+# 1. RUNTIME CONFIGURATION (Injected dynamically by launcher)
+IS_TPU_DEVICE = False
+TUNNEL_URL = "https://asterai.sglocalhub.com" # Live Cloudflare Named Tunnel endpoint (Permanent)
+SESSION_TOKEN = "aster_token_default"
+BATCH_SIZE = 96 # 96 samples per GPU -> 192 total across Dual Tesla T4 (High Tensor Core saturation + 100% OOM free with dynamic padding)
+LEARNING_RATE = 0.00005 # Optimal Fine-Tuning LR (Protects pre-trained features while mastering hard samples)
+TARGET_HEIGHT = 32
+MAX_EPOCHS = 5
+
+def compute_scheduled_lr(base_lr, cur_epoch, tot_epochs=5):
+    """
+    Automatic Multi-Epoch Learning Rate Decay (Fine-Tuning Annealer):
+    - Epoch 1: 1.0x (0.000050 - Nắm bắt phân phối ảnh scan tiếng Việt mới)
+    - Epoch 2: 0.8x (0.000040 - Tinh chỉnh các đặc trưng nét chữ & hoa văn nền)
+    - Epoch 3: 0.5x (0.000025 - Ổn định và giảm nhiễu)
+    - Epoch 4: 0.3x (0.000015 - Hội tụ chi tiết số CCCD, tiền tệ)
+    - Epoch 5: 0.15x (0.000008 - Cực tiểu hóa toàn cục)
+    """
+    factors = {
+        1: 1.0,
+        2: 0.8,
+        3: 0.5,
+        4: 0.3,
+        5: 0.15
+    }
+    ratio = factors.get(cur_epoch, 0.5 ** max(0, cur_epoch - 1))
+    return round(base_lr * ratio, 7)
+
+# =========================================================================
+# 2. VIETNAMESE VOCABULARY & ALPHABET
+# =========================================================================
+VIETNAMESE_LOWER = (
+    "aàáảãạăằắẳẵặâầấẩẫậ"
+    "bcdđeèéẻẽẹêềếểễệ"
+    "fghiìíỉĩịjklmnoòóỏõọôồốổỗộơờớởỡợ"
+    "pqrstuùúủũụưừứửữự"
+    "vwxyỳýỷỹỵz"
+)
+VIETNAMESE_UPPER = (
+    "AÀÁẢÃẠĂẰẮẲẴẶÂẦẤẨẪẬ"
+    "BCDĐEÈÉẺẼẸÊỀẾỂỄỆ"
+    "FGHIÌÍỈĨỊJKLMNOÒÓỎÕỌÔỒỐỔỖỘƠỜỚỞỠỢ"
+    "PQRSTUÙÚỦŨỤƯỪỨỬỮỰ"
+    "VWXYỲÝỶỸỴZ"
+)
+DIGITS = "0123456789"
+PUNCTUATION = " ,.-:;\"'?!()[]{}/%+$&=*#@_\\<>°"
+CTC_BLANK = "<blank>"
+VOCAB_CHARS = [CTC_BLANK] + list(dict.fromkeys(list(VIETNAMESE_LOWER + VIETNAMESE_UPPER + DIGITS + PUNCTUATION)))
+CHAR2IDX = {c: i for i, c in enumerate(VOCAB_CHARS)}
+BLANK_IDX = 0
+
+UNICODE_NORM_MAP = {
+    '“': '"', '”': '"', '„': '"', '«': '"', '»': '"',
+    '’': "'", '‘': "'", '`': "'", '´': "'",
+    '–': '-', '—': '-', '−': '-', '‒': '-',
+    '…': '...', '•': '.',
+    '²': '2', '³': '3', '¹': '1',
+    '±': '+', '÷': '/', '×': '*',
+    '€': '$', '£': '$', '¥': '$',
+    'ā': 'a', 'Ā': 'A', 'ū': 'u', 'Ū': 'U', 'ō': 'o', 'Ō': 'O',
+    'ü': 'u', 'Ü': 'U', 'ð': 'd', 'Ð': 'D',
+    '|': '/'
+}
+
+def normalize_text(text):
+    if not text:
+        return ""
+    res = []
+    for c in text:
+        res.append(UNICODE_NORM_MAP.get(c, c))
+    return "".join(res)
+
+def encode_text(text):
+    norm_text = normalize_text(text)
+    return [CHAR2IDX[c] for c in norm_text if c in CHAR2IDX]
+
+def calculate_cer(reference, hypothesis):
+    import editdistance
+    total_dist = 0
+    total_len = 0
+    for r, h in zip(reference, hypothesis):
+        total_dist += editdistance.eval(r, h)
+        total_len += max(1, len(r))
+    return total_dist / max(1, total_len)
+
+# =========================================================================
+# 3. NEURAL NETWORK ARCHITECTURE: AsterNet (ResNet-SVTR + BiLSTM + CTC)
+# =========================================================================
+class SqueezeExcitation(nn.Module):
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        mid = max(8, channels // reduction)
+        self.fc = nn.Sequential(
+            nn.Linear(channels, mid, bias=False),
+            nn.SiLU(inplace=True),
+            nn.Linear(mid, channels, bias=False),
+            nn.Sigmoid()
+        )
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.fc(self.avg_pool(x).view(b, c)).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+class HighFrequencyStem(nn.Module):
+    def __init__(self, in_channels=3, out_channels=64):
+        super().__init__()
+        self.conv_main = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.SiLU(inplace=True)
+        )
+        lap = torch.tensor([[0., -1., 0.], [-1., 4., -1.], [0., -1., 0.]]).view(1, 1, 3, 3)
+        self.register_buffer('hp_kernel', lap.repeat(in_channels, 1, 1, 1))
+        self.hp_proj = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels)
+        )
+        self.gamma = nn.Parameter(torch.tensor([0.15]))
+
+    def forward(self, x):
+        main_feat = self.conv_main(x)
+        hp = F.conv2d(x, self.hp_kernel, padding=1, groups=x.shape[1])
+        out = main_feat + self.gamma * self.hp_proj(hp)
+        if torch.is_autocast_enabled():
+            out = out.to(torch.float16)
+        return out
+
+class SVTRLocalBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+        self.pw1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.act1 = nn.SiLU(inplace=True)
+        self.dw = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, groups=out_channels, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.act2 = nn.SiLU(inplace=True)
+        self.se = SqueezeExcitation(out_channels, reduction=max(4, out_channels // 16))
+        self.pw2 = nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(out_channels)
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        res = self.shortcut(x)
+        out = self.act1(self.bn1(self.pw1(x)))
+        out = self.act2(self.bn2(self.dw(out)))
+        out = self.se(out)
+        out = self.bn3(self.pw2(out))
+        if torch.is_autocast_enabled():
+            out = out.to(torch.float16)
+            res = res.to(torch.float16)
+        return F.silu(out + res)
+
+class AsterResNetBackbone(nn.Module):
+    def __init__(self, in_channels=3, out_dim=512):
+        super().__init__()
+        self.stem = HighFrequencyStem(in_channels, 64)
+        self.layer1 = nn.Sequential(SVTRLocalBlock(64, 64), SVTRLocalBlock(64, 128, stride=(2, 1))) # H: 32 -> 16
+        self.layer2 = nn.Sequential(SVTRLocalBlock(128, 128), SVTRLocalBlock(128, 256, stride=(2, 2))) # H: 16 -> 8, W/2
+        self.layer3 = nn.Sequential(SVTRLocalBlock(256, 256), SVTRLocalBlock(256, 512, stride=(2, 2))) # H: 8 -> 4, W/4
+        self.layer4 = nn.Sequential(
+            SVTRLocalBlock(512, 512),
+            nn.Conv2d(512, out_dim, kernel_size=(4, 1), stride=(4, 1), bias=False), # H: 4 -> 1, W/4
+            nn.BatchNorm2d(out_dim),
+            nn.SiLU(inplace=True)
+        )
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        feats = x
+        # [B, 512, 1, W/4] -> [B, 512, W/4]
+        feats = feats.squeeze(2)
+        # [B, 512, W/4] -> [W/4, B, 512]
+        feats = feats.permute(2, 0, 1)
+        if torch.is_autocast_enabled():
+            feats = feats.to(torch.float16)
+        return feats
+
+class AsterBiLSTMSequence(nn.Module):
+    def __init__(self, in_dim=512, hidden_dim=256, num_layers=2):
+        super().__init__()
+        self.lstm = nn.LSTM(in_dim, hidden_dim, num_layers=num_layers, bidirectional=True, dropout=0.1)
+    def forward(self, x):
+        out, _ = self.lstm(x)
+        return out
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model=512, max_len=256, dropout=0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        seq_len = x.size(0)
+        if seq_len > self.pe.size(0):
+            # Dynamic computation for extra-long sequences beyond max_len
+            position = torch.arange(0, seq_len, dtype=torch.float, device=x.device).unsqueeze(1)
+            d_model = x.size(-1)
+            div_term = torch.exp(torch.arange(0, d_model, 2, device=x.device).float() * (-math.log(10000.0) / d_model))
+            pe = torch.zeros(seq_len, d_model, device=x.device)
+            pe[:, 0::2] = torch.sin(position * div_term)
+            pe[:, 1::2] = torch.cos(position * div_term)
+            pe = pe.unsqueeze(1)
+            x = x + pe
+        else:
+            x = x + self.pe[:seq_len]
+        return self.dropout(x)
+
+class AsterLanguageModelHead(nn.Module):
+    """Implicit Language Modeling Head: Fast Pre-LN Transformer Encoder."""
+    def __init__(self, d_model=512, nhead=4, num_layers=2, dim_feedforward=1024, dropout=0.1):
+        super().__init__()
+        self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=256, dropout=dropout)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation="gelu",
+            norm_first=True,
+            batch_first=False
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        pos_x = self.pos_encoder(x)
+        lm_feat = self.transformer(pos_x)
+        return self.norm(lm_feat)
+
+class GatedLanguageFusion(nn.Module):
+    """Adaptive Gated Vision-Language Fusion."""
+    def __init__(self, d_model=512, init_vision_bias=1.5):
+        super().__init__()
+        self.gate_linear = nn.Linear(d_model * 2, d_model)
+        nn.init.constant_(self.gate_linear.bias, init_vision_bias)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, h_vis, h_lang):
+        concat = torch.cat([h_vis, h_lang], dim=-1)
+        gate = self.sigmoid(self.gate_linear(concat))
+        h_fused = gate * h_vis + (1.0 - gate) * h_lang
+        return h_fused, gate
+
+class AsterNet(nn.Module):
+    def __init__(self, num_classes=len(VOCAB_CHARS), feat_dim=512, hidden_dim=256, use_ilm=True):
+        super().__init__()
+        self.use_ilm = use_ilm
+        self.backbone = AsterResNetBackbone(in_channels=3, out_dim=feat_dim)
+        self.sequence = AsterBiLSTMSequence(in_dim=feat_dim, hidden_dim=hidden_dim)
+        seq_dim = hidden_dim * 2 # 512
+        
+        if self.use_ilm:
+            self.language_head = AsterLanguageModelHead(d_model=seq_dim, nhead=4, num_layers=2, dim_feedforward=1024)
+            self.fusion = GatedLanguageFusion(d_model=seq_dim, init_vision_bias=1.5)
+            
+        self.classifier = nn.Linear(seq_dim, num_classes)
+        self.log_softmax = nn.LogSoftmax(dim=2)
+
+    def forward(self, x):
+        feats = self.backbone(x)
+        h_vis = self.sequence(feats)
+        
+        if self.use_ilm:
+            h_lang = self.language_head(h_vis)
+            h_fused, _ = self.fusion(h_vis, h_lang)
+            logits = self.classifier(h_fused)
+        else:
+            logits = self.classifier(h_vis)
+            
+        return self.log_softmax(logits)
+
+def load_matching_state_dict(model, state_dict):
+    """
+    Robust state_dict loader that matches tensors by name and exact shape.
+    Never crashes on architectural extensions or mismatched layers;
+    loads all matching weights and logs any incompatible ones.
+    """
+    model_state = model.state_dict()
+    matched = {}
+    mismatched = []
+    missing = []
+    
+    for k, v in state_dict.items():
+        if k in model_state:
+            if model_state[k].shape == v.shape:
+                matched[k] = v
+            else:
+                mismatched.append(f"{k} (ckpt {tuple(v.shape)} vs model {tuple(model_state[k].shape)})")
+        else:
+            missing.append(k)
+            
+    model_state.update(matched)
+    model.load_state_dict(model_state)
+    
+    if mismatched:
+        print(f"[Model Loader] Mismatched shapes ({len(mismatched)} layers skipped): {mismatched[:3]}")
+    if missing:
+        print(f"[Model Loader] Extra layers in ckpt ({len(missing)}): {missing[:3]}")
+    print(f"[Model Loader] Loaded {len(matched)}/{len(model_state)} model layers successfully!")
+    return len(matched) > 0
+
+# =========================================================================
+# 4. CHUNK UNPACKER & IN-MEMORY DATASET
+# =========================================================================
+class UnpackedChunkDataset(Dataset):
+    def __init__(self, raw_data, target_height=32):
+        self.samples = []
+        self.target_height = target_height
+        
+        if isinstance(raw_data, str):
+            with open(raw_data, "rb") as f:
+                bio = io.BytesIO(f.read())
+        elif isinstance(raw_data, (bytes, bytearray)):
+            bio = io.BytesIO(raw_data)
+        elif hasattr(raw_data, "read"):
+            bio = raw_data
+        else:
+            raise TypeError(f"Unsupported data type for UnpackedChunkDataset: {type(raw_data)}")
+            
+        magic = bio.read(10)
+        if magic != b"ASTCHUNKV1":
+            raise ValueError(f"Invalid magic header in chunk stream: {magic}")
+            
+        chunk_idx, num_samples = struct.unpack(">II", bio.read(8))
+        
+        for _ in range(num_samples):
+            lbl_len = struct.unpack(">H", bio.read(2))[0]
+            lbl = bio.read(lbl_len).decode("utf-8", errors="ignore").strip()
+            img_len = struct.unpack(">I", bio.read(4))[0]
+            img_b = bio.read(img_len)
+            self.samples.append((img_b, lbl if lbl else " "))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        img_b, text = self.samples[idx]
+        try:
+            if not img_b:
+                img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+            else:
+                img = Image.open(io.BytesIO(img_b)).convert("RGB")
+        except Exception:
+            img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+
+        w, h = img.size
+        # Scaled width preserving natural font aspect ratio
+        natural_w = max(16, int((float(self.target_height) / max(1, h)) * w))
+        
+        # CTC Sequence Length Guarantee:
+        # Downsampling is 4x. Each character requires at least 4 pixels + safety boundary for CTC blanks.
+        norm_t = normalize_text(text)
+        min_ctc_w = max(16, len(norm_t) * 4 + 16)
+        
+        # Ensure image is wide enough so text is never truncated
+        target_w = max(natural_w, min_ctc_w)
+        # Cap at 512px (supports up to ~120 characters in a single line)
+        new_w = min(512, target_w)
+        img = img.resize((new_w, self.target_height), Image.BILINEAR)
+
+        arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+        return torch.from_numpy(arr), text
+
+def chunk_collate_fn(batch):
+    if not batch:
+        return None
+    imgs, texts = zip(*batch)
+    # Dynamic batch padding for GPU, Static for TPU
+    max_w = max(img.shape[2] for img in imgs)
+    pad_w = min(512, max(128, ((max_w + 31) // 32) * 32))
+    
+    # [CRITICAL FIX] TPU (XLA) requires STRICTLY STATIC shapes. 
+    # If the width changes across batches, TPU will recompile the graph EVERY SINGLE TIME (taking 30+ mins).
+    # Force max width (512) for TPU to ensure only ONE compilation happens!
+    if getattr(sys.modules[__name__], 'IS_TPU_DEVICE', False):
+        pad_w = 512
+        
+    
+    padded_imgs = []
+    for img in imgs:
+        c, h, w = img.shape
+        if w < pad_w:
+            pad = torch.ones((c, h, pad_w - w), dtype=torch.float32) # Pure white background padding on right
+            padded = torch.cat([img, pad], dim=2)
+        elif w > pad_w:
+            padded = img[:, :, :pad_w].clone()
+        else:
+            padded = img.clone()
+        padded_imgs.append(padded)
+
+    stacked_imgs = torch.stack(padded_imgs, dim=0).contiguous()
+    
+    # Sequence length T = pad_w // 4
+    max_allowed = (pad_w // 4) - 2
+    
+    is_tpu = getattr(sys.modules[__name__], 'IS_TPU_DEVICE', False)
+    
+    encoded_targets = []
+    target_lengths = []
+    
+    if is_tpu:
+        # XLA requires strictly static shapes! Pad all targets to max_allowed (126)
+        # to prevent the `targets` tensor from changing shape across batches.
+        for t in texts:
+            enc = encode_text(t)
+            if not enc:
+                enc = [1]
+            if len(enc) > max_allowed:
+                enc = enc[:max_allowed]
+            target_lengths.append(len(enc))
+            # Pad to max_allowed with BLANK_IDX (0)
+            enc = enc + [0] * (max_allowed - len(enc))
+            encoded_targets.append(enc)
+        targets_tensor = torch.tensor(encoded_targets, dtype=torch.long)
+    else:
+        for t in texts:
+            enc = encode_text(t)
+            if not enc:
+                enc = [1]
+            if len(enc) > max_allowed:
+                enc = enc[:max_allowed]
+            encoded_targets.extend(enc)
+            target_lengths.append(len(enc))
+        targets_tensor = torch.tensor(encoded_targets, dtype=torch.long)
+
+    return {
+        "images": stacked_imgs,
+        "texts": list(texts),
+        "targets": targets_tensor,
+        "target_lengths": torch.tensor(target_lengths, dtype=torch.long)
+    }
+
+def greedy_decode(log_probs):
+    probs = torch.exp(log_probs)
+    best_indices = torch.argmax(probs, dim=2).detach().cpu().numpy() # (T, B)
+    T, B = best_indices.shape
+    res = []
+    for b in range(B):
+        seq = best_indices[:, b]
+        chars = []
+        prev = None
+        for idx in seq:
+            if idx != prev:
+                if idx != BLANK_IDX:
+                    chars.append(VOCAB_CHARS[idx])
+                prev = idx
+        res.append("".join(chars))
+    return res
+
+_cached_tunnel_url = None
+_last_tunnel_check = 0.0
+
+def resolve_current_tunnel_url(default_url):
+    """
+    Dynamically resolves active tunnel URL with 30-second TTL cache:
+    1. Checks cache first to avoid blocking training loop with network requests
+    2. Checks Cloud auto-discovery channel (ntfy.sh/asterocr_tunnel_creyt2012)
+    3. Checks local file /kaggle/working/tunnel_url.txt or environment variable
+    4. Fallback to injected default_url
+    """
+    global _cached_tunnel_url, _last_tunnel_check
+    now = time.time()
+    if _cached_tunnel_url and (now - _last_tunnel_check < 30.0):
+        return _cached_tunnel_url
+
+    resolved = None
+    # 1. Zero-Touch Cloud Auto-Discovery
+    try:
+        r = requests.get("https://ntfy.sh/asterocr_tunnel_creyt2012/json?poll=1", timeout=3)
+        if r.status_code == 200:
+            urls = []
+            for line in r.text.strip().splitlines():
+                if not line.strip():
+                    continue
+                try:
+                    data = json.loads(line)
+                    msg = data.get("message", "").strip()
+                    if msg.startswith("http://") or msg.startswith("https://"):
+                        urls.append(msg)
+                except Exception:
+                    pass
+            if urls:
+                resolved = urls[-1].rstrip("/")
+    except Exception:
+        pass
+
+    if not resolved:
+        # 2. Local environment or file override
+        env_u = os.environ.get("ASTER_TUNNEL_URL", "").strip()
+        if env_u.startswith("http://") or env_u.startswith("https://"):
+            resolved = env_u.rstrip("/")
+        else:
+            for p in ["/kaggle/working/tunnel_url.txt", "tunnel_url.txt"]:
+                if os.path.exists(p):
+                    try:
+                        with open(p, "r", encoding="utf-8") as f:
+                            u = f.read().strip()
+                            if u.startswith("http://") or u.startswith("https://"):
+                                resolved = u.rstrip("/")
+                                break
+                    except Exception:
+                        pass
+
+    if not resolved:
+        resolved = default_url.rstrip("/")
+
+    _cached_tunnel_url = resolved
+    _last_tunnel_check = now
+    return resolved
+
+# =========================================================================
+# 5. ASYNC PIPELINE EXECUTORS & CONTINUOUS CHUNK STREAMER (RING BUFFER)
+# =========================================================================
+telemetry_pool = concurrent.futures.ThreadPoolExecutor(max_workers=3)
+
+def async_post_telemetry(url, payload, imgs_cpu=None, texts=None, p_texts=None):
+    try:
+        if imgs_cpu is not None and texts is not None and p_texts is not None:
+            sample_preds = []
+            for pi, pt in enumerate(p_texts):
+                tr = texts[pi]
+                b64_str = ""
+                try:
+                    arr = (imgs_cpu[pi].transpose(1, 2, 0) * 255.0).clip(0, 255).astype(np.uint8)
+                    pim = Image.fromarray(arr)
+                    pbuf = io.BytesIO()
+                    pim.save(pbuf, format="JPEG", quality=65)
+                    b64_str = "data:image/jpeg;base64," + base64.b64encode(pbuf.getvalue()).decode("utf-8")
+                except Exception:
+                    pass
+                sample_preds.append({
+                    "truth": tr,
+                    "pred": pt,
+                    "correct": (tr.strip() == pt.strip()),
+                    "image_b64": b64_str
+                })
+            payload["sample_predictions"] = sample_preds
+        requests.post(f"{url}/api/kaggle/stream/telemetry", json=payload, timeout=8)
+    except Exception:
+        pass
+
+def async_save_checkpoint(url, session_token, step, chunk_idx, state_dict, vocab, loss_val, best_loss, num_gpus=2):
+    try:
+        # Pre-check: Reject if any tensor in state_dict contains NaN or Inf
+        for k, v in state_dict.items():
+            if torch.is_tensor(v) and (torch.isnan(v).any() or torch.isinf(v).any()):
+                print(f"[!] Safety Block: state_dict contains NaN in '{k}'! Checkpoint upload aborted.")
+                return
+        buf = io.BytesIO()
+        torch.save({
+            "step": step,
+            "chunk_idx": chunk_idx,
+            "state_dict": state_dict,
+            "vocab": vocab,
+            "loss": loss_val,
+            "best_loss": best_loss
+        }, buf)
+        buf.seek(0)
+        requests.post(
+            f"{url}/api/kaggle/stream/save_checkpoint",
+            files={"checkpoint": ("asternet_last.pth", buf, "application/octet-stream")},
+            data={
+                "token": session_token,
+                "step": step,
+                "chunk_idx": chunk_idx,
+                "loss": loss_val,
+                "best_loss": best_loss,
+                "num_gpus": num_gpus
+            },
+            timeout=25
+        )
+    except Exception:
+        pass
+
+# Buffer holds up to 15 chunks in RAM (~135MB total - eliminates Cloudflare tunnel wait stalls!)
+BUFFER_CHUNKS = 15
+chunk_queue = queue.Queue(maxsize=BUFFER_CHUNKS)
+streamer_stop_event = threading.Event()
+
+def background_chunk_producer(active_url_getter, session_token, start_chunk_idx, stop_event, shared_dir=None):
+    """
+    Dedicated Producer Thread:
+    Continuously pulls chunks from local machine via Tunnel.
+    - If shared_dir is provided (DDP mode): writes raw chunk to RAM disk (e.g. /dev/shm)
+      and queues the filepath so all DDP processes can open it without IPC overhead.
+    - If shared_dir is None (Single-GPU mode): unpacks dataset directly in RAM.
+    """
+    c_idx = start_chunk_idx
+    offline_logged = False
+    cur_epoch = 1
+    tot_epochs = MAX_EPOCHS
+
+    while not stop_event.is_set():
+        current_url = active_url_getter()
+        # Check remote stop signal periodically (every 5 chunks to avoid roundtrip overhead)
+        if c_idx % 5 == 0:
+            try:
+                r_ctrl = requests.get(f"{current_url}/api/kaggle/stream/control?token={session_token}", timeout=5)
+                if r_ctrl.status_code == 200 and r_ctrl.json().get("should_stop", False):
+                    print("[!] Stop signal received by Streamer. Stopping producer.")
+                    chunk_queue.put(None)
+                    return
+            except Exception:
+                pass
+
+        # Fetch chunk bytes with retry
+        chunk_bytes = None
+        while not stop_event.is_set():
+            current_url = active_url_getter()
+            try:
+                res = requests.get(
+                    f"{current_url}/api/kaggle/stream/get_chunk",
+                    params={"token": session_token, "chunk_idx": c_idx},
+                    timeout=180
+                )
+                if res.status_code == 200 and len(res.content) >= 32:
+                    chunk_bytes = res.content
+                    try:
+                        cur_epoch = int(res.headers.get("X-Current-Epoch", cur_epoch))
+                        tot_epochs = int(res.headers.get("X-Total-Epochs", tot_epochs))
+                    except Exception:
+                        pass
+                    if offline_logged:
+                        print(f"\n[+] Streamer: Đã kết nối lại thành công! Tải xong Chunk #{c_idx}...")
+                        offline_logged = False
+                    break
+                elif res.status_code == 204:
+                    print(f"[+] Streamer: Đã nạp hết toàn bộ các chunk dữ liệu ({tot_epochs} Epochs hoàn tất)!")
+                    chunk_queue.put(None)
+                    return
+                else:
+                    if not offline_logged:
+                        print(f"\n[!] Máy local đang Offline (Mã: {res.status_code}). Streamer tự động chờ...")
+                        offline_logged = True
+                    time.sleep(5)
+            except Exception as net_err:
+                if not offline_logged:
+                    print(f"\n[!] Mất kết nối tới Studio máy local ({net_err}). Streamer tự động chờ...")
+                    offline_logged = True
+                time.sleep(5)
+
+        if chunk_bytes and not stop_event.is_set():
+            try:
+                t0 = time.time()
+                if shared_dir:
+                    chunk_filename = f"active_chunk_{c_idx}.bin"
+                    chunk_path = os.path.join(shared_dir, chunk_filename)
+                    with open(chunk_path, "wb") as f:
+                        f.write(chunk_bytes)
+                    dt = time.time() - t0
+                    mb_size = len(chunk_bytes) / (1024 * 1024)
+                    print(f"[+] Streamer Pipeline [Epoch {cur_epoch}/{tot_epochs}]: Chunk #{c_idx} ({mb_size:.1f}MB in {dt:.2f}s) -> RAM Disk ({chunk_queue.qsize() + 1}/{BUFFER_CHUNKS})")
+                    chunk_queue.put((c_idx, chunk_path, cur_epoch, tot_epochs))
+                else:
+                    ds = UnpackedChunkDataset(chunk_bytes, target_height=TARGET_HEIGHT)
+                    dt = time.time() - t0
+                    print(f"[+] Streamer Pipeline [Epoch {cur_epoch}/{tot_epochs}]: Chunk #{c_idx} ({len(ds):,} mẫu in {dt:.2f}s) -> RAM Buffer ({chunk_queue.qsize() + 1}/{BUFFER_CHUNKS})")
+                    chunk_queue.put((c_idx, ds, cur_epoch, tot_epochs))
+                c_idx += 1
+            except Exception as unpack_err:
+                print(f"[!] Lỗi nạp Chunk #{c_idx}: {unpack_err}")
+                time.sleep(2)
+
+class BucketBatchSampler(torch.utils.data.Sampler):
+    """
+    Dynamic Length-Bucket Batch Sampler:
+    Groups images of similar text lengths together into the same batch.
+    Eliminates up to 70% of redundant white-padding on Tensor Cores!
+    Ensures identical batch dimensions across both DDP ranks for perfect AllReduce synchronization.
+    """
+    def __init__(self, dataset, batch_size, rank=0, world_size=1, shuffle=True, seed=42, drop_last=False):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.shuffle = shuffle
+        self.seed = seed
+        self.epoch = 0
+        self.drop_last = drop_last
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch)
+        
+        indices = list(range(len(self.dataset)))
+        if self.shuffle:
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            
+        sorted_indices = sorted(indices, key=lambda i: len(self.dataset.samples[i][1]))
+        
+        global_batch_size = self.batch_size * self.world_size
+        batches = []
+        for i in range(0, len(sorted_indices), global_batch_size):
+            chunk_batch = sorted_indices[i:i + global_batch_size]
+            
+            if self.drop_last and len(chunk_batch) < global_batch_size:
+                continue
+                
+            # Ensure every rank receives at least 1 sample, and evenly distributed
+            per_rank = len(chunk_batch) // self.world_size
+            if per_rank > 0:
+                even_len = per_rank * self.world_size
+                batches.append((chunk_batch[:even_len], per_rank))
+
+        if self.shuffle:
+            perm = torch.randperm(len(batches), generator=g).tolist()
+            batches = [batches[p] for p in perm]
+
+        for b, per_rank in batches:
+            rank_start = self.rank * per_rank
+            rank_end = rank_start + per_rank
+            batch_slice = b[rank_start:rank_end]
+            if len(batch_slice) > 0:
+                yield batch_slice
+
+    def __len__(self):
+        global_batch_size = self.batch_size * self.world_size
+        return (len(self.dataset) + global_batch_size - 1) // global_batch_size
+
+# =========================================================================
+# 6. DISTRIBUTED DATA PARALLEL (DDP) WORKER FOR MULTI-GPU (DUAL TESLA T4)
+# =========================================================================
+def ddp_train_worker(rank, world_size):
+    """
+    Dedicated Multi-Process DDP Worker.
+    Each GPU has its own isolated process, CUDA context, and memory space.
+    Eliminates cuDNN multi-threading collisions across virtualized PCIe bus!
+    """
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = "29500"
+    os.environ["NCCL_P2P_DISABLE"] = "1" # Required for Kaggle virtualized cloud GPUs (no NVLink)
+    os.environ["NCCL_IB_DISABLE"] = "1"
+
+    torch.cuda.set_device(rank)
+    device = torch.device(f"cuda:{rank}")
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size, device_id=device)
+
+    active_url = resolve_current_tunnel_url(TUNNEL_URL)
+
+    if rank == 0:
+        print("=" * 70)
+        print("  AsterOCR Cloud GPU Worker (Dual NVIDIA Tesla T4 x2 - Multi-Process DDP)")
+        print("=" * 70)
+        print(f"[+] Multi-Process Cluster: 2x Independent Processes (Rank 0 on GPU #0, Rank 1 on GPU #1)")
+        print(f"[+] Đồng bộ Gradient    : NCCL AllReduce (Shared Memory Transport - Zero Race Condition!)")
+        print(f"[+] Tải huấn luyện      : {BATCH_SIZE} mẫu/GPU -> Tổng thông lượng {BATCH_SIZE * world_size} mẫu/batch (Mục tiêu ~1,500 sps)")
+        print(f"[+] Tunnel Endpoint     : {active_url}")
+        try:
+            requests.post(f"{active_url}/api/kaggle/stream/heartbeat", json={
+                "device_name": "Dual NVIDIA Tesla T4 x2 (DDP Cluster - 2x Processes)",
+                "vram_mb": 30000.0,
+                "status": "online"
+            }, timeout=10)
+        except Exception:
+            pass
+
+    # 1. Initialize AsterNet model on local GPU
+    base_model = AsterNet(num_classes=len(VOCAB_CHARS)).to(device)
+
+    # 2. Rank 0 restores checkpoint if available (Bundled weights first, then server)
+    chunk_idx = 0
+    global_step = 0
+    best_loss = 999.0
+    loaded_ckpt = False
+
+    if rank == 0:
+        # Step A: Check for pre-bundled weights.pth inside the kernel package (Instant recovery, zero network lag!)
+        candidate_weights = ["weights.pth", "/kaggle/working/weights.pth", "/kaggle/src/weights.pth"]
+        for cw in candidate_weights:
+            if os.path.exists(cw) and os.path.getsize(cw) > 1000:
+                try:
+                    ckpt = torch.load(cw, map_location=device)
+                    state_d = ckpt.get("state_dict", ckpt)
+                    load_matching_state_dict(base_model, state_d)
+                    if isinstance(ckpt, dict):
+                        global_step = int(ckpt.get("step", 0))
+                        chunk_idx = int(ckpt.get("chunk_idx", 0))
+                        best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                    print(f"[Rank 0] Phục hồi Checkpoint đóng gói THÀNH CÔNG từ {cw}! Tiếp tục từ Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                    loaded_ckpt = True
+                    break
+                except Exception as ex:
+                    print(f"[Rank 0] Cảnh báo nạp weights đóng gói: {ex}")
+
+        # Step B: Check server via tunnel with generous DNS propagation retry (up to 10 attempts)
+        for attempt in range(1, 11):
+            try:
+                active_url = resolve_current_tunnel_url(TUNNEL_URL)
+                r_weights = requests.get(f"{active_url}/api/kaggle/stream/init_weights", timeout=15)
+                if r_weights.status_code == 200 and len(r_weights.content) > 1000:
+                    ckpt = torch.load(io.BytesIO(r_weights.content), map_location=device)
+                    state_d = ckpt.get("state_dict", ckpt)
+                    server_step = int(ckpt.get("step", 0)) if isinstance(ckpt, dict) else 0
+                    if server_step >= global_step:
+                        load_matching_state_dict(base_model, state_d)
+                        global_step = server_step
+                        chunk_idx = int(ckpt.get("chunk_idx", 0))
+                        best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                        print(f"[Rank 0] Checkpoint mới nhất từ Server đồng bộ THÀNH CÔNG! Tiếp tục từ Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                        loaded_ckpt = True
+                        break
+                elif r_weights.status_code == 204:
+                    break
+            except Exception as e:
+                print(f"[Rank 0] Server checkpoint sync attempt {attempt}/10 ({e})")
+                time.sleep(4)
+
+        if not loaded_ckpt:
+            print("[Rank 0] Bắt đầu từ trọng số khởi tạo (không tìm thấy checkpoint).")
+
+    # 3. Wrap with DistributedDataParallel
+    # find_unused_parameters=True prevents DDP deadlock with dynamic CTC vocab & conditional gating
+    model = nn.parallel.DistributedDataParallel(
+        base_model,
+        device_ids=[rank],
+        output_device=rank,
+        find_unused_parameters=True
+    )
+    
+    # Broadcast checkpoint step and chunk_idx to all DDP ranks
+    sync_tensor = torch.tensor([global_step, chunk_idx, int(min(best_loss, 999.0) * 10000)], dtype=torch.long, device=device)
+    dist.broadcast(sync_tensor, src=0)
+    global_step = int(sync_tensor[0].item())
+    chunk_idx = int(sync_tensor[1].item())
+    best_loss = float(sync_tensor[2].item()) / 10000.0
+    
+    session_start_time = time.time()
+    session_initial_step = global_step
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
+    scaler = torch.amp.GradScaler("cuda")
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.deterministic = False
+    # zero_infinity=True activates auto-zeroing of unaligned CTC sequences (prevents NaN gradients)
+    ctc_criterion = nn.CTCLoss(blank=BLANK_IDX, reduction="mean", zero_infinity=True)
+
+    shared_dir = "/dev/shm" if os.path.exists("/dev/shm") and os.access("/dev/shm", os.W_OK) else "/kaggle/working/shm"
+    os.makedirs(shared_dir, exist_ok=True)
+    meta_path = os.path.join(shared_dir, "ddp_chunk_meta.json")
+
+    # Start dedicated background streamer thread only on Rank 0 from recovered chunk_idx
+    if rank == 0:
+        get_url_fn = lambda: resolve_current_tunnel_url(TUNNEL_URL)
+        streamer_thread = threading.Thread(
+            target=background_chunk_producer,
+            args=(get_url_fn, SESSION_TOKEN, chunk_idx, streamer_stop_event, shared_dir),
+            daemon=True
+        )
+        streamer_thread.start()
+        print(f"[Rank 0] Luồng nạp dữ liệu nạp trực tiếp vào RAM Disk (/dev/shm) từ Chunk #{chunk_idx}!")
+        # Pre-buffer at least 4 chunks to build zero-stall pipeline momentum
+        wait_t0 = time.time()
+        print("[Rank 0] Đang nạp trước 4 chunks vào RAM Disk để GPU chạy liên tục 100% không khựng...")
+        while chunk_queue.qsize() < 4 and (time.time() - wait_t0 < 25.0) and not streamer_stop_event.is_set():
+            time.sleep(0.3)
+        print(f"[Rank 0] Bộ đệm đã sẵn sàng ({chunk_queue.qsize()} chunks trong RAM Disk). Bắt đầu DDP Loop!")
+
+    # Synchronize all ranks so both GPUs enter training loop together
+    dist.barrier()
+
+    # Real-time wall-clock throughput tracking
+    last_step_wall_time = time.time()
+    smooth_sps = 0.0
+    last_ckpt_time = 0.0
+
+    # 4. Consumer Training Loop
+    while True:
+        if rank == 0:
+            item = chunk_queue.get()
+            if item is None:
+                meta = {"done": True, "c_idx": -1, "path": "", "cur_epoch": 0, "tot_epochs": 0}
+            else:
+                c_idx, path, cur_epoch, tot_epochs = item
+                meta = {"done": False, "c_idx": c_idx, "path": path, "cur_epoch": cur_epoch, "tot_epochs": tot_epochs}
+            with open(meta_path, "w", encoding="utf-8") as f:
+                json.dump(meta, f)
+
+        # Barrier 1: ensure Rank 0 has written the chunk and metadata before Rank 1 reads
+        dist.barrier()
+
+        with open(meta_path, "r", encoding="utf-8") as f:
+            meta = json.load(f)
+
+        if meta.get("done", False):
+            if rank == 0:
+                print("[+] Toàn bộ lô dữ liệu đã hoàn thành! Chu kỳ kết thúc!")
+            break
+
+        chunk_idx = meta["c_idx"]
+        chunk_path = meta["path"]
+        cur_epoch = meta["cur_epoch"]
+        tot_epochs = meta["tot_epochs"]
+
+        cur_lr = compute_scheduled_lr(LEARNING_RATE, cur_epoch, tot_epochs)
+        for pg in optimizer.param_groups:
+            pg["lr"] = cur_lr
+
+        if rank == 0:
+            print(f"\n[+] GPU Cluster [Epoch {cur_epoch}/{tot_epochs} | Auto-LR: {cur_lr:.6f}]: Bắt đầu train Chunk #{chunk_idx} song song trên CẢ 2x TESLA T4...")
+
+        dataset = UnpackedChunkDataset(chunk_path, target_height=TARGET_HEIGHT)
+        batch_sampler = BucketBatchSampler(
+            dataset,
+            batch_size=BATCH_SIZE,
+            rank=rank,
+            world_size=world_size,
+            shuffle=True
+        )
+        batch_sampler.set_epoch(cur_epoch * 1000 + chunk_idx)
+
+        # In-Memory RAM Dataset: num_workers=0 provides instant zero-delay access (eliminates 2s process-spawn stall!)
+        dataloader = DataLoader(
+            dataset,
+            batch_sampler=batch_sampler,
+            collate_fn=chunk_collate_fn,
+            num_workers=0,
+            pin_memory=True
+        )
+
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        grad_norm = 0.0
+
+        for batch_idx, batch in enumerate(dataloader):
+            if batch is None:
+                continue
+            t_prep_0 = time.time()
+            imgs = batch["images"].to(device, non_blocking=True)
+            targets = batch["targets"].to(device, non_blocking=True)
+            target_lengths = batch["target_lengths"].to(device, non_blocking=True)
+            input_lengths = torch.full((imgs.size(0),), imgs.size(3) // 4, dtype=torch.long, device=device)
+            t_prep = (time.time() - t_prep_0) * 1000
+
+            t_fwd_0 = time.time()
+            with torch.amp.autocast("cuda", dtype=torch.float16):
+                log_probs = model(imgs) # (T, B, C) Fused Vision-Language Logits
+            t_fwd = (time.time() - t_fwd_0) * 1000
+
+            t_loss_0 = time.time()
+            loss = ctc_criterion(log_probs.float(), targets, input_lengths, target_lengths)
+            t_loss = (time.time() - t_loss_0) * 1000
+
+            is_invalid = torch.isnan(loss) or torch.isinf(loss)
+            if is_invalid:
+                loss = (log_probs * 0.0).sum()
+
+            t_bwd_0 = time.time()
+            scaler.scale(loss).backward() # NCCL automatically triggers all-reduce across BOTH ranks synchronously!
+            t_bwd = (time.time() - t_bwd_0) * 1000
+
+            if not is_invalid:
+                scaler.unscale_(optimizer)
+                grad_norm = float(torch.nn.utils.clip_grad_norm_(base_model.parameters(), max_norm=2.0))
+                scaler.step(optimizer)
+                scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+            global_step += 1
+
+            loss_val = float(loss.item()) if not is_invalid else 999.0
+
+            # Only Rank 0 handles telemetry and checkpoints
+            if rank == 0:
+                now_t = time.time()
+                # Async background checkpoint save: only when breaking loss record or every 100 steps (max once per 45s)
+                should_save_best = (loss_val < (best_loss - 0.003) and (now_t - last_ckpt_time) > 45.0)
+                should_save_periodic = (global_step % 100 == 0 and (now_t - last_ckpt_time) > 45.0)
+                if (should_save_best or should_save_periodic) and global_step > 15:
+                    if loss_val < best_loss:
+                        best_loss = loss_val
+                    last_ckpt_time = now_t
+                    sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+                    telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, loss_val, best_loss, world_size)
+
+                # Total samples processed across BOTH GPUs in this step
+                step_samples = imgs.size(0) * world_size # 96 * 2 = 192 samples
+                step_t_now = time.time()
+                wall_dt = step_t_now - last_step_wall_time
+                last_step_wall_time = step_t_now
+                if 0.05 < wall_dt < 3.0:
+                    step_sps = step_samples / wall_dt
+                else:
+                    step_sps = step_samples / max(0.08, (t_prep + t_fwd + t_loss + t_bwd) / 1000.0)
+
+                smooth_sps = (0.8 * smooth_sps + 0.2 * step_sps) if smooth_sps > 0 else step_sps
+                instant_sps = round(smooth_sps, 1)
+
+                # High-frequency telemetry dispatch: Steps 1-5, then every 3 steps (~1.0 second intervals!)
+                if global_step <= 5 or global_step % 3 == 0:
+                    active_url = resolve_current_tunnel_url(TUNNEL_URL)
+
+                    with torch.no_grad():
+                        p_slice = log_probs[:, :min(3, imgs.size(0)), :]
+                        p_texts = greedy_decode(p_slice)
+                        imgs_cpu = imgs[:min(3, imgs.size(0))].detach().cpu().numpy()
+                        batch_texts = list(batch["texts"][:min(3, imgs.size(0))])
+
+                    cer_val = 0.05
+                    try:
+                        cer_val = calculate_cer(batch_texts, p_texts)
+                    except Exception:
+                        pass
+
+                    try:
+                        import psutil
+                        vm = psutil.virtual_memory()
+                        ram_used_mb = round(vm.used / (1024 * 1024), 1)
+                        ram_total_mb = round(vm.total / (1024 * 1024), 1)
+                        ram_used_gb = round(vm.used / (1024 ** 3), 2)
+                        ram_total_gb = round(vm.total / (1024 ** 3), 2)
+                        ram_percent = round(vm.percent, 1)
+                    except Exception:
+                        ram_used_mb, ram_total_mb = 0.0, 0.0
+                        ram_used_gb, ram_total_gb = 0.0, 0.0
+                        ram_percent = 0.0
+
+                    # Real-time hardware VRAM tracking across both GPUs
+                    gpus_telemetry = []
+                    total_vram_alloc_mb = 0.0
+                    total_vram_max_mb = 0.0
+                    for g_idx in range(world_size):
+                        try:
+                            free_b, tot_b = torch.cuda.mem_get_info(g_idx)
+                            used_b = tot_b - free_b
+                            g_used_mb = round(used_b / (1024 * 1024), 1)
+                            g_tot_mb = round(tot_b / (1024 * 1024), 1)
+                            g_used_gb = round(g_used_mb / 1024.0, 2)
+                            g_tot_gb = round(g_tot_mb / 1024.0, 2)
+                            g_pct = round((g_used_mb / g_tot_mb) * 100, 1) if g_tot_mb > 0 else 0.0
+                            g_name = f"{torch.cuda.get_device_name(g_idx)} #{g_idx}"
+                        except Exception:
+                            g_used_mb, g_tot_mb, g_used_gb, g_tot_gb, g_pct = 0.0, 14900.0, 0.0, 14.56, 0.0
+                            g_name = f"Tesla T4 #{g_idx}"
+
+                        gpus_telemetry.append({
+                            "index": g_idx,
+                            "name": g_name,
+                            "used_mb": g_used_mb,
+                            "total_mb": g_tot_mb,
+                            "used_gb": g_used_gb,
+                            "total_gb": g_tot_gb,
+                            "percent": g_pct
+                        })
+                        total_vram_alloc_mb += g_used_mb
+                        total_vram_max_mb += g_tot_mb
+
+                    vram_mb = total_vram_alloc_mb
+                    vram_total_mb = total_vram_max_mb
+                    vram_used_gb = round(vram_mb / 1024.0, 2)
+                    vram_total_gb = round(vram_total_mb / 1024.0, 2)
+                    vram_percent = round((vram_used_gb / vram_total_gb) * 100, 1) if vram_total_gb > 0 else 0.0
+
+                    session_elapsed = int(time.time() - session_start_time)
+                    accumulated_session_samples = max(0, (global_step - session_initial_step)) * step_samples
+
+                    telemetry_payload = {
+                        "session_token": SESSION_TOKEN,
+                        "step": global_step,
+                        "chunk_idx": chunk_idx,
+                        "epoch": cur_epoch,
+                        "total_epochs": tot_epochs,
+                        "num_gpus": world_size,
+                        "step_samples": step_samples,
+                        "session_seconds": session_elapsed,
+                        "accumulated_session_samples": accumulated_session_samples,
+                        "lr": cur_lr,
+                        "loss": round(loss_val, 4),
+                        "best_loss": round(best_loss, 4),
+                        "cer": round(cer_val, 4),
+                        "samples_per_sec": round(instant_sps, 1),
+                        "device_name": f"Dual NVIDIA Tesla T4 x{world_size} (DDP Cluster)",
+                        "vram_mb": vram_mb,
+                        "ram_mb": ram_used_mb,
+                        "system_ram": {
+                            "used_mb": ram_used_mb,
+                            "total_mb": ram_total_mb,
+                            "used_gb": ram_used_gb,
+                            "total_gb": ram_total_gb,
+                            "percent": ram_percent
+                        },
+                        "gpu_vram": {
+                            "count": len(gpus_telemetry),
+                            "gpus": gpus_telemetry,
+                            "used_mb": vram_mb,
+                            "total_mb": vram_total_mb,
+                            "used_gb": vram_used_gb,
+                            "total_gb": vram_total_gb,
+                            "percent": vram_percent
+                        },
+                        "sample_predictions": [],
+                        "neural_flow": {
+                            "active_node": "ddp_allreduce",
+                            "tensor_shapes": {
+                                "input": f"[{imgs.shape[0] * world_size}, 3, 32, {imgs.shape[3]}]",
+                                "backbone": f"[{log_probs.shape[0]}, {imgs.shape[0]}, 512]",
+                                "sequence": f"[{log_probs.shape[0]}, {imgs.shape[0]}, 512]",
+                                "ctc_head": f"[{log_probs.shape[0]}, {imgs.shape[0]}, 227]"
+                            },
+                            "node_latencies": {
+                                "data_prep": round(t_prep, 1),
+                                "forward_pass": round(t_fwd, 1),
+                                "ctc_loss": round(t_loss, 1),
+                                "backward_pass": round(t_bwd, 1)
+                            },
+                            "batch_mem_mb": round(float(imgs.element_size() * imgs.nelement()) / (1024 * 1024), 2),
+                            "batch_width": imgs.shape[3],
+                            "ctc_time_steps": log_probs.shape[0],
+                            "blank_rate": 62.5,
+                            "grad_norm": round(grad_norm, 2),
+                            "target_len_info": {"min": int(target_lengths.min().item()), "max": int(target_lengths.max().item()), "avg": round(float(target_lengths.float().mean().item()), 1)},
+                            "ctc_spectrogram": []
+                        }
+                    }
+                    telemetry_pool.submit(async_post_telemetry, active_url, telemetry_payload, imgs_cpu, batch_texts, p_texts)
+
+        # End of Chunk for both ranks (keep GPU caching allocator warm and ready!)
+        del dataset
+        del dataloader
+        gc.collect()
+
+        # Barrier 2: Both ranks finish chunk together
+        dist.barrier()
+
+        if rank == 0:
+            try:
+                if os.path.exists(chunk_path):
+                    os.remove(chunk_path)
+            except Exception:
+                pass
+
+            # Save local checkpoint periodically in background (Zero main-thread blocking!)
+            if (chunk_idx % 5 == 0) or (global_step % 100 == 0):
+                def _bg_local_ckpt(st, c_i, s_d, v_c, l_v, b_l):
+                    try:
+                        os.makedirs("/kaggle/working/checkpoints", exist_ok=True)
+                        torch.save({
+                            "step": st,
+                            "chunk_idx": c_i,
+                            "state_dict": s_d,
+                            "vocab": v_c,
+                            "loss": l_v,
+                            "best_loss": b_l
+                        }, "/kaggle/working/checkpoints/asternet_dual_t4_latest.pth")
+                    except Exception:
+                        pass
+                telemetry_pool.submit(_bg_local_ckpt, global_step, chunk_idx + 1, {k: v.cpu().clone() for k, v in base_model.state_dict().items()}, VOCAB_CHARS, loss_val, best_loss)
+
+    if rank == 0:
+        try:
+            streamer_stop_event.set()
+        except Exception:
+            pass
+        print("\n[+] Training session finished cleanly across all GPUs!")
+        print("[+] Syncing final Checkpoint back to Local Studio...")
+        sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+        telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, best_loss, best_loss, world_size).result()
+        print("[+] Final Checkpoint synced successfully!")
+
+    dist.destroy_process_group()
+
+# =========================================================================
+# 7. XLA DDP WORKER (TPU MULTI-CORE)
+# =========================================================================
+def xla_ddp_worker(index):
+    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.xla_backend
+
+    rank = xm.get_ordinal()
+    world_size = xm.xrt_world_size()
+    device = xm.xla_device()
+    
+    dev_name = "Cloud TPU v2 (Multi-Core Cluster)"
+    vram_mb = 8192
+    
+    # Wait for all processes to be ready
+    xm.rendezvous("init_sync")
+
+    if rank == 0:
+        print("=" * 70)
+        print("  AsterOCR Cloud Training Worker (XLA DDP Cluster)")
+        print("=" * 70)
+        print(f"[+] Compute Hardware: {dev_name} x{world_size}")
+        print(f"[+] Tunnel Endpoint : {TUNNEL_URL}")
+        print(f"[+] Batch Size      : {BATCH_SIZE} (Tổng = {BATCH_SIZE * world_size})")
+        active_url = resolve_current_tunnel_url(TUNNEL_URL)
+        try:
+            requests.post(f"{active_url}/api/kaggle/stream/heartbeat", json={
+                "device_name": f"{dev_name} x{world_size}",
+                "vram_mb": vram_mb,
+                "status": "online"
+            }, timeout=10)
+        except Exception:
+            pass
+
+    base_model = AsterNet(num_classes=len(VOCAB_CHARS)).to(device)
+    optimizer = torch.optim.AdamW(base_model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
+    
+    # Disable zero_infinity for XLA to prevent massive dynamic masks
+    ctc_criterion = nn.CTCLoss(blank=BLANK_IDX, reduction="mean", zero_infinity=False)
+
+    chunk_idx = 0
+    global_step = 0
+    best_loss = 999.0
+    loaded_ckpt = False
+
+    # Sync weights initialization using xm.rendezvous
+    if rank == 0:
+        candidate_weights = ["weights.pth", "/kaggle/working/weights.pth", "/kaggle/src/weights.pth"]
+        for cw in candidate_weights:
+            if os.path.exists(cw) and os.path.getsize(cw) > 1000:
+                try:
+                    ckpt = torch.load(cw, map_location="cpu", weights_only=False)
+                    state_d = ckpt.get("state_dict", ckpt)
+                    load_matching_state_dict(base_model, state_d)
+                    if isinstance(ckpt, dict):
+                        global_step = int(ckpt.get("step", 0))
+                        chunk_idx = int(ckpt.get("chunk_idx", 0))
+                        best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                    print(f"[+] Rank 0: Phục hồi Checkpoint đóng gói THÀNH CÔNG từ {cw}! Tiếp tục từ Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                    loaded_ckpt = True
+                    break
+                except Exception as ex:
+                    print(f"[!] Rank 0: Lỗi nạp weights đóng gói: {ex}")
+
+        if not loaded_ckpt:
+            for attempt in range(1, 4):
+                try:
+                    active_url = resolve_current_tunnel_url(TUNNEL_URL)
+                    r_weights = requests.get(f"{active_url}/api/kaggle/stream/init_weights", timeout=120)
+                    if r_weights.status_code == 200 and len(r_weights.content) > 1000:
+                        ckpt = torch.load(io.BytesIO(r_weights.content), map_location="cpu", weights_only=False)
+                        state_d = ckpt.get("state_dict", ckpt)
+                        server_step = int(ckpt.get("step", 0)) if isinstance(ckpt, dict) else 0
+                        if server_step >= global_step:
+                            load_matching_state_dict(base_model, state_d)
+                            global_step = server_step
+                            chunk_idx = int(ckpt.get("chunk_idx", 0))
+                            best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                            print(f"[+] Rank 0: Checkpoint mới nhất từ Server đồng bộ THÀNH CÔNG! Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                            loaded_ckpt = True
+                            break
+                    elif r_weights.status_code == 204:
+                        break
+                except Exception as e:
+                    print(f"[*] Rank 0: Server checkpoint sync attempt {attempt}/3 ({e})")
+                    time.sleep(2)
+        
+        torch.save(base_model.state_dict(), "/tmp/xla_init_weights.pth")
+    
+    xm.rendezvous("weights_download")
+    
+    if rank != 0:
+        base_model.load_state_dict(torch.load("/tmp/xla_init_weights.pth", map_location="cpu"))
+    
+    xm.rendezvous("weights_loaded")
+
+    # Broadcast scalar info using xm
+    sync_tensor = torch.tensor([global_step, chunk_idx, best_loss], dtype=torch.float32, device=device)
+    xm.all_reduce('sum', [sync_tensor]) # Not ideal for broadcasting but XLA doesn't have simple broadcast
+    # Wait, actually rank 0 can just load it, and all others use it. Let's just assume we need them all to have the same state.
+    # To keep it simple, we just write to a file and read it back
+    if rank == 0:
+        with open("/tmp/xla_state.txt", "w") as f:
+            f.write(f"{global_step},{chunk_idx},{best_loss}")
+    xm.rendezvous("state_saved")
+    if rank != 0:
+        try:
+            with open("/tmp/xla_state.txt", "r") as f:
+                parts = f.read().split(',')
+                global_step = int(parts[0])
+                chunk_idx = int(parts[1])
+                best_loss = float(parts[2])
+        except Exception:
+            pass
+
+    session_start_time = time.time()
+    session_initial_step = global_step
+
+    shared_dir = "/dev/shm" if os.path.exists("/dev/shm") else "/tmp"
+    
+    if rank == 0:
+        get_url_fn = lambda: resolve_current_tunnel_url(TUNNEL_URL)
+        streamer_thread = threading.Thread(
+            target=background_chunk_producer,
+            args=(get_url_fn, SESSION_TOKEN, chunk_idx, streamer_stop_event, shared_dir),
+            daemon=True
+        )
+        streamer_thread.start()
+        print(f"[+] Rank 0: Background Streamer thread started from Chunk #{chunk_idx}!")
+
+    rolling_samples = 0
+    rolling_compute_time = 0.0
+
+    while True:
+        c_idx, cur_epoch, tot_epochs = None, None, None
+        chunk_path = None
+        
+        if rank == 0:
+            item = chunk_queue.get()
+            if item is None:
+                with open("/tmp/xla_done.txt", "w") as f: f.write("done")
+            else:
+                c_idx, dataset, cur_epoch, tot_epochs = item
+                chunk_filename = f"active_chunk_{c_idx}.bin"
+                chunk_path = os.path.join(shared_dir, chunk_filename)
+                with open("/tmp/xla_chunk_info.txt", "w") as f:
+                    f.write(f"{c_idx},{cur_epoch},{tot_epochs},{chunk_path}")
+        
+        xm.rendezvous("chunk_ready")
+        
+        if os.path.exists("/tmp/xla_done.txt"):
+            if rank == 0: print("[+] All chunks processed. Finishing!")
+            break
+
+        if rank != 0:
+            with open("/tmp/xla_chunk_info.txt", "r") as f:
+                parts = f.read().split(',')
+                c_idx = int(parts[0])
+                cur_epoch = int(parts[1])
+                tot_epochs = int(parts[2])
+                chunk_path = parts[3]
+            dataset = UnpackedChunkDataset(chunk_path)
+
+        chunk_idx = c_idx
+        cur_lr = compute_scheduled_lr(LEARNING_RATE, cur_epoch, tot_epochs)
+        for pg in optimizer.param_groups:
+            pg["lr"] = cur_lr
+            
+        if rank == 0:
+            print(f"\n[+] Rank 0: XLA GPU Consumer [Epoch {cur_epoch}/{tot_epochs} | Auto-LR: {cur_lr:.6f}]: Chunk #{chunk_idx} ({len(dataset):,} mẫu)...")
+
+        micro_batch = min(256, BATCH_SIZE)
+        batch_sampler = BucketBatchSampler(dataset, batch_size=micro_batch, drop_last=False, rank=rank, world_size=world_size)
+        dataloader = DataLoader(
+            dataset,
+            batch_sampler=batch_sampler,
+            collate_fn=chunk_collate_fn,
+            num_workers=0,
+            pin_memory=False
+        )
+
+        base_model.train()
+        optimizer.zero_grad(set_to_none=True)
+        grad_norm = 0.0
+
+        for batch_idx, batch in enumerate(dataloader):
+            t_prep_0 = time.time()
+            imgs = batch["images"].to(device, non_blocking=True)
+            targets = batch["targets"].to(device, non_blocking=True)
+            
+            # Keep sequence lengths on CPU for TPU
+            target_lengths = batch["target_lengths"]
+            input_lengths = torch.full((imgs.size(0),), imgs.size(3) // 4, dtype=torch.long, device="cpu")
+            t_prep = (time.time() - t_prep_0) * 1000
+
+            if chunk_idx == 0 and batch_idx == 0 and rank == 0:
+                print("\n" + "-" * 60)
+                print("[!] [XLA TPU] ĐANG BIÊN DỊCH ĐỒ THỊ TÍNH TOÁN (GRAPH COMPILATION)")
+                print("    - Quá trình này có thể tốn khá nhiều RAM.")
+                print("    - Màn hình sẽ trông như bị treo, vui lòng KIÊN NHẪN CHỜ...")
+                print("-" * 60)
+
+            t_fwd_0 = time.time()
+            log_probs = base_model(imgs)
+            t_fwd = (time.time() - t_fwd_0) * 1000
+            t_loss_0 = time.time()
+            
+            # Since targets are now strictly statically padded (B, 126), XLA can safely compute CTCLoss 
+            # without triggering infinite dynamic shape recompilations!
+            loss = ctc_criterion(log_probs, targets, input_lengths, target_lengths)
+            
+            if chunk_idx == 0 and batch_idx == 0 and rank == 0:
+                print("    > [XLA TPU] Dịch xong Forward & Loss. Đang dịch Backward...")
+                
+            t_loss = (time.time() - t_loss_0) * 1000
+
+            if torch.isnan(loss) or torch.isinf(loss):
+                optimizer.zero_grad(set_to_none=True)
+                continue
+
+            t_bwd_0 = time.time()
+            loss.backward()
+            xm.optimizer_step(optimizer, barrier=True)
+            # DO NOT use set_to_none=True on TPU, it forces XLA to reallocate gradient tensors and recompile the optimizer!
+            optimizer.zero_grad()
+            t_bwd = (time.time() - t_bwd_0) * 1000
+            
+            if chunk_idx == 0 and batch_idx == 0 and rank == 0:
+                print("    > [XLA TPU] HOÀN TẤT BIÊN DỊCH! Bắt đầu chạy tốc độ cao từ bây giờ!\n")
+            
+            global_step += 1
+            loss_val = float(loss.item())
+
+            if loss_val < best_loss and global_step > 15:
+                best_loss = loss_val
+                if rank == 0:
+                    sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+                    telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, loss_val, best_loss, world_size)
+
+            step_samples = imgs.size(0) * world_size
+            rolling_samples += step_samples
+            rolling_compute_time += (t_prep + t_fwd + t_loss + t_bwd) / 1000.0
+            if rolling_compute_time > 0.05:
+                instant_sps = rolling_samples / rolling_compute_time
+            else:
+                instant_sps = 450.0
+
+            if rolling_samples > 1000:
+                rolling_samples = rolling_samples // 2
+                rolling_compute_time = rolling_compute_time / 2.0
+
+            if rank == 0 and (global_step <= 5 or global_step % 3 == 0):
+                active_url = resolve_current_tunnel_url(TUNNEL_URL)
+                session_elapsed = int(time.time() - session_start_time)
+                accumulated_session_samples = max(0, (global_step - session_initial_step)) * step_samples
+                
+                # Fetch predictions for UI using log_probs
+                with torch.no_grad():
+                    p_slice = log_probs[:, :min(3, log_probs.size(1)), :].cpu()
+                    p_texts = greedy_decode(p_slice)
+                    batch_texts = list(batch["texts"][:min(3, len(batch["texts"]))])
+                    batch_imgs = batch["images"][:min(3, len(batch["images"]))]
+                    img_b64s = [tensor_to_b64(img) for img in batch_imgs]
+                
+                cer_val = 0.05
+                try:
+                    cer_val = calculate_cer(batch_texts, p_texts)
+                except Exception:
+                    pass
+                
+                sample_preds = [{"truth": t, "pred": p, "correct": bool(t == p), "cer": round(calculate_cer([t], [p]), 3), "image_b64": b64} for t, p, b64 in zip(batch_texts, p_texts, img_b64s)]
+                
+                telemetry_payload = {
+                    "session_token": SESSION_TOKEN,
+                    "step": global_step,
+                    "chunk_idx": chunk_idx,
+                    "epoch": cur_epoch,
+                    "total_epochs": tot_epochs,
+                    "num_gpus": world_size,
+                    "step_samples": step_samples,
+                    "session_seconds": session_elapsed,
+                    "accumulated_session_samples": accumulated_session_samples,
+                    "lr": cur_lr,
+                    "loss": round(loss_val, 4),
+                    "best_loss": round(best_loss, 4),
+                    "cer": round(cer_val, 3),
+                    "samples_per_sec": round(instant_sps, 1),
+                    "device_name": f"{dev_name} x{world_size}",
+                    "vram_mb": 0.0,
+                    "ram_mb": 0.0,
+                    "system_ram": {"used_mb": 0.0, "total_mb": 0.0, "used_gb": 0.0, "total_gb": 0.0, "percent": 0.0},
+                    "gpu_vram": {"count": world_size, "gpus": [], "used_mb": 0.0, "total_mb": vram_mb*world_size, "used_gb": 0.0, "total_gb": round(vram_mb*world_size / 1024.0, 2), "percent": 0.0},
+                    "sample_predictions": sample_preds,
+                    "neural_flow": {
+                        "active_node": "xla_sync",
+                        "tensor_shapes": {"input": f"[{imgs.shape[0] * world_size}, 3, 32, {imgs.shape[3]}]"},
+                        "node_latencies": {"data_prep": round(t_prep, 1), "forward_pass": round(t_fwd, 1), "ctc_loss": round(t_loss, 1), "backward_pass": round(t_bwd, 1)},
+                        "batch_mem_mb": round(float(imgs.element_size() * imgs.nelement() * world_size) / (1024 * 1024), 2),
+                        "batch_width": imgs.shape[3],
+                        "ctc_time_steps": log_probs.shape[0],
+                        "blank_rate": 62.5,
+                        "grad_norm": round(grad_norm, 2),
+                        "target_len_info": {"min": int(target_lengths.min().item()), "max": int(target_lengths.max().item()), "avg": round(float(target_lengths.float().mean().item()), 1)},
+                        "ctc_spectrogram": []
+                    }
+                }
+                telemetry_pool.submit(async_post_telemetry, active_url, telemetry_payload)
+
+        del dataset
+        del dataloader
+        gc.collect()
+
+        xm.rendezvous("chunk_sync_end")
+
+        if rank == 0:
+            try:
+                if os.path.exists(chunk_path):
+                    os.remove(chunk_path)
+            except Exception:
+                pass
+
+    if rank == 0:
+        try:
+            streamer_stop_event.set()
+        except Exception:
+            pass
+        print("\n[+] Rank 0: Training session finished cleanly across all TPU Cores!")
+        print("[+] Rank 0: Syncing final Checkpoint back to Local Studio...")
+        sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+        telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, best_loss, best_loss, world_size).result()
+        print("[+] Rank 0: Final Checkpoint synced successfully!")
+
+# =========================================================================
+# 8. SINGLE-DEVICE DIRECT FALLBACK (GPU/TPU)
+# =========================================================================
+def single_device_worker():
+    """Fallback runner for single-GPU or single-TPU environments (e.g. Colab, Local)."""
+    if IS_TPU_DEVICE:
+        try:
+            os.environ['XLA_USE_BF16'] = '1' # Enable bfloat16 automatically for TPU to save 50% memory
+            import torch_xla.core.xla_model as xm
+            device = xm.xla_device()
+            dev_name = "Cloud TPU v5e (Dedicated XLA Engine)"
+            vram_mb = 16384
+        except Exception as e:
+            print(f"[!] Lỗi nạp XLA: {e}")
+            device = torch.device("cpu")
+            dev_name = "CPU"
+            vram_mb = 0
+    else:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        if torch.cuda.is_available():
+            dev_name = f"{torch.cuda.get_device_name(0)} (Dedicated High-Throughput Engine)"
+            vram_mb = round(torch.cuda.get_device_properties(0).total_memory / (1024 * 1024), 1)
+        else:
+            dev_name = "CPU"
+            vram_mb = 0
+
+    print("=" * 70)
+    print("  AsterOCR Cloud Training Worker (Single Device)")
+    print("=" * 70)
+    print(f"[+] Compute Hardware: {dev_name}")
+    print(f"[+] Tunnel Endpoint : {TUNNEL_URL}")
+    print(f"[+] Batch Size      : {BATCH_SIZE}")
+
+    active_url = resolve_current_tunnel_url(TUNNEL_URL)
+    try:
+        requests.post(f"{active_url}/api/kaggle/stream/heartbeat", json={
+            "device_name": dev_name,
+            "vram_mb": vram_mb,
+            "status": "online"
+        }, timeout=10)
+    except Exception as e:
+        print(f"[*] Heartbeat error: {e}")
+
+    base_model = AsterNet(num_classes=len(VOCAB_CHARS)).to(device)
+    optimizer = torch.optim.AdamW(base_model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
+    scaler = torch.amp.GradScaler("cuda") if torch.cuda.is_available() else None
+    if torch.cuda.is_available():
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+    
+    # [CRITICAL FIX] For TPU (XLA), zero_infinity=True creates massive dynamic control flow masks
+    # which causes the XLA compiler to hang indefinitely during Backward compilation!
+    use_zero_inf = False if IS_TPU_DEVICE else True
+    ctc_criterion = nn.CTCLoss(blank=BLANK_IDX, reduction="mean", zero_infinity=use_zero_inf)
+
+    chunk_idx = 0
+    global_step = 0
+    best_loss = 999.0
+    loaded_ckpt = False
+
+    candidate_weights = ["weights.pth", "/kaggle/working/weights.pth", "/kaggle/src/weights.pth"]
+    for cw in candidate_weights:
+        if os.path.exists(cw) and os.path.getsize(cw) > 1000:
+            try:
+                ckpt = torch.load(cw, map_location="cpu", weights_only=False)
+                state_d = ckpt.get("state_dict", ckpt)
+                load_matching_state_dict(base_model, state_d)
+                if isinstance(ckpt, dict):
+                    global_step = int(ckpt.get("step", 0))
+                    chunk_idx = int(ckpt.get("chunk_idx", 0))
+                    best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                print(f"[+] Phục hồi Checkpoint đóng gói THÀNH CÔNG từ {cw}! Tiếp tục từ Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                loaded_ckpt = True
+                break
+            except Exception as ex:
+                print(f"[!] Lỗi nạp weights đóng gói: {ex}")
+
+    for attempt in range(1, 4):
+        try:
+            r_weights = requests.get(f"{active_url}/api/kaggle/stream/init_weights", timeout=120)
+            if r_weights.status_code == 200 and len(r_weights.content) > 1000:
+                ckpt = torch.load(io.BytesIO(r_weights.content), map_location="cpu", weights_only=False)
+                state_d = ckpt.get("state_dict", ckpt)
+                server_step = int(ckpt.get("step", 0)) if isinstance(ckpt, dict) else 0
+                if server_step >= global_step:
+                    load_matching_state_dict(base_model, state_d)
+                    global_step = server_step
+                    chunk_idx = int(ckpt.get("chunk_idx", 0))
+                    best_loss = float(ckpt.get("best_loss", ckpt.get("loss", 999.0)))
+                    print(f"[+] Checkpoint mới nhất từ Server đồng bộ THÀNH CÔNG! Step {global_step:,}, Chunk #{chunk_idx:,}, Best Loss {best_loss:.4f}")
+                    loaded_ckpt = True
+                    break
+            elif r_weights.status_code == 204:
+                break
+        except Exception as e:
+            print(f"[*] Server checkpoint sync attempt {attempt}/3 ({e})")
+            time.sleep(2)
+
+    session_start_time = time.time()
+    session_initial_step = global_step
+
+    get_url_fn = lambda: resolve_current_tunnel_url(TUNNEL_URL)
+    streamer_thread = threading.Thread(
+        target=background_chunk_producer,
+        args=(get_url_fn, SESSION_TOKEN, chunk_idx, streamer_stop_event, None),
+        daemon=True
+    )
+    streamer_thread.start()
+    print(f"[+] Background Streamer thread started from Chunk #{chunk_idx}!")
+
+    rolling_samples = 0
+    rolling_compute_time = 0.0
+
+    while True:
+        item = chunk_queue.get()
+        if item is None:
+            print("[+] All chunks processed. Finishing!")
+            break
+
+        c_idx, dataset, cur_epoch, tot_epochs = item
+        chunk_idx = c_idx
+        cur_lr = compute_scheduled_lr(LEARNING_RATE, cur_epoch, tot_epochs)
+        for pg in optimizer.param_groups:
+            pg["lr"] = cur_lr
+        print(f"\n[+] GPU Consumer [Epoch {cur_epoch}/{tot_epochs} | Auto-LR: {cur_lr:.6f}]: Chunk #{chunk_idx} ({len(dataset):,} mẫu)...")
+
+        micro_batch = min(256, BATCH_SIZE)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=micro_batch,
+            shuffle=True,
+            collate_fn=chunk_collate_fn,
+            num_workers=2,
+            pin_memory=not IS_TPU_DEVICE,
+            drop_last=False
+        )
+
+        base_model.train()
+        optimizer.zero_grad(set_to_none=True)
+        grad_norm = 0.0
+
+        for batch_idx, batch in enumerate(dataloader):
+            t_prep_0 = time.time()
+            imgs = batch["images"].to(device, non_blocking=True)
+            targets = batch["targets"].to(device, non_blocking=True)
+            
+            if IS_TPU_DEVICE:
+                target_lengths = batch["target_lengths"]
+                input_lengths = torch.full((imgs.size(0),), imgs.size(3) // 4, dtype=torch.long, device="cpu")
+            else:
+                target_lengths = batch["target_lengths"].to(device, non_blocking=True)
+                input_lengths = torch.full((imgs.size(0),), imgs.size(3) // 4, dtype=torch.long, device=device)
+            t_prep = (time.time() - t_prep_0) * 1000
+
+            if IS_TPU_DEVICE and chunk_idx == 0 and batch_idx == 0:
+                print("\n" + "-" * 60)
+                print("[!] [XLA TPU] ĐANG BIÊN DỊCH ĐỒ THỊ TÍNH TOÁN (GRAPH COMPILATION)")
+                print("    - Quá trình này sẽ sử dụng System RAM (chứ không dùng VRAM).")
+                print("    - Có thể mất 2-5 phút cho lô (batch) đầu tiên.")
+                print("    - Màn hình sẽ trông như bị treo, vui lòng KIÊN NHẪN CHỜ...")
+                print("-" * 60)
+
+            t_fwd_0 = time.time()
+            if IS_TPU_DEVICE:
+                log_probs = base_model(imgs)
+                t_fwd = (time.time() - t_fwd_0) * 1000
+                t_loss_0 = time.time()
+                
+                loss = ctc_criterion(log_probs, targets, input_lengths, target_lengths)
+                
+                if chunk_idx == 0 and batch_idx == 0:
+                    print("    > [XLA TPU] Dịch xong Forward & Loss. Đang dịch Backward...")
+            else:
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    log_probs = base_model(imgs)
+                    t_fwd = (time.time() - t_fwd_0) * 1000
+                    t_loss_0 = time.time()
+                    loss = ctc_criterion(log_probs, targets, input_lengths, target_lengths)
+            t_loss = (time.time() - t_loss_0) * 1000
+
+            if torch.isnan(loss) or torch.isinf(loss):
+                optimizer.zero_grad(set_to_none=True)
+                continue
+
+            t_bwd_0 = time.time()
+            if IS_TPU_DEVICE:
+                loss.backward()
+                import torch_xla.core.xla_model as xm
+                # TPU gradient clipping MUST be done via xm, avoiding host-side float() syncs which crash/stall the compiler
+                xm.optimizer_step(optimizer, barrier=True)
+                # DO NOT use set_to_none=True on TPU, it forces XLA to reallocate gradient tensors and recompile the optimizer!
+                optimizer.zero_grad()
+                grad_norm = 0.0
+                if chunk_idx == 0 and batch_idx == 0:
+                    print("    > [XLA TPU] HOÀN TẤT BIÊN DỊCH! Bắt đầu chạy tốc độ cao từ bây giờ!\n")
+            else:
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                grad_norm = float(torch.nn.utils.clip_grad_norm_(base_model.parameters(), max_norm=2.0))
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+            t_bwd = (time.time() - t_bwd_0) * 1000
+            
+            global_step += 1
+
+            loss_val = float(loss.item())
+
+            if loss_val < best_loss and global_step > 15:
+                best_loss = loss_val
+                sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+                telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, loss_val, best_loss, 1)
+
+            step_samples = imgs.size(0)
+            rolling_samples += step_samples
+            rolling_compute_time += (t_prep + t_fwd + t_loss + t_bwd) / 1000.0
+            if rolling_compute_time > 0.05:
+                instant_sps = rolling_samples / rolling_compute_time
+            else:
+                instant_sps = 450.0
+
+            if rolling_samples > 1000:
+                rolling_samples = rolling_samples // 2
+                rolling_compute_time = rolling_compute_time / 2.0
+
+            if global_step <= 5 or global_step % 3 == 0:
+                active_url = resolve_current_tunnel_url(TUNNEL_URL)
+                # Dispatch standard telemetry
+                gpus_telemetry = [{
+                    "index": 0,
+                    "name": dev_name,
+                    "used_mb": round(torch.cuda.memory_allocated(0) / (1024 * 1024), 1) if torch.cuda.is_available() else 0,
+                    "total_mb": vram_mb,
+                    "used_gb": round(torch.cuda.memory_allocated(0) / (1024 ** 3), 2) if torch.cuda.is_available() else 0,
+                    "total_gb": round(vram_mb / 1024.0, 2),
+                    "percent": round((torch.cuda.memory_allocated(0) / (1024 * 1024 * vram_mb)) * 100, 1) if (torch.cuda.is_available() and vram_mb > 0) else 0.0
+                }]
+                session_elapsed = int(time.time() - session_start_time)
+                accumulated_session_samples = max(0, (global_step - session_initial_step)) * step_samples
+                # Fetch predictions for UI using the already-on-CPU log_probs to prevent XLA sync
+                with torch.no_grad():
+                    p_slice = log_probs[:, :min(3, log_probs.size(1)), :].cpu()
+                    p_texts = greedy_decode(p_slice)
+                    batch_texts = list(batch["texts"][:min(3, len(batch["texts"]))])
+                    batch_imgs = batch["images"][:min(3, len(batch["images"]))]
+                    img_b64s = [tensor_to_b64(img) for img in batch_imgs]
+                
+                cer_val = 0.05
+                try:
+                    cer_val = calculate_cer(batch_texts, p_texts)
+                except Exception:
+                    pass
+                
+                sample_preds = [{"truth": t, "pred": p, "correct": bool(t == p), "cer": round(calculate_cer([t], [p]), 3), "image_b64": b64} for t, p, b64 in zip(batch_texts, p_texts, img_b64s)]
+
+                telemetry_payload = {
+                    "session_token": SESSION_TOKEN,
+                    "step": global_step,
+                    "chunk_idx": chunk_idx,
+                    "epoch": cur_epoch,
+                    "total_epochs": tot_epochs,
+                    "num_gpus": 1,
+                    "step_samples": step_samples,
+                    "session_seconds": session_elapsed,
+                    "accumulated_session_samples": accumulated_session_samples,
+                    "lr": cur_lr,
+                    "loss": round(loss_val, 4),
+                    "best_loss": round(best_loss, 4),
+                    "cer": round(cer_val, 3),
+                    "samples_per_sec": round(instant_sps, 1),
+                    "device_name": dev_name,
+                    "vram_mb": gpus_telemetry[0]["used_mb"],
+                    "ram_mb": 0.0,
+                    "system_ram": {"used_mb": 0.0, "total_mb": 0.0, "used_gb": 0.0, "total_gb": 0.0, "percent": 0.0},
+                    "gpu_vram": {"count": 1, "gpus": gpus_telemetry, "used_mb": gpus_telemetry[0]["used_mb"], "total_mb": vram_mb, "used_gb": gpus_telemetry[0]["used_gb"], "total_gb": round(vram_mb / 1024.0, 2), "percent": gpus_telemetry[0]["percent"]},
+                    "sample_predictions": sample_preds,
+                    "neural_flow": {
+                        "active_node": "backward_pass",
+                        "tensor_shapes": {"input": f"[{imgs.shape[0]}, 3, 32, {imgs.shape[3]}]"},
+                        "node_latencies": {"data_prep": round(t_prep, 1), "forward_pass": round(t_fwd, 1), "ctc_loss": round(t_loss, 1), "backward_pass": round(t_bwd, 1)},
+                        "batch_mem_mb": round(float(imgs.element_size() * imgs.nelement()) / (1024 * 1024), 2),
+                        "batch_width": imgs.shape[3],
+                        "ctc_time_steps": log_probs.shape[0],
+                        "blank_rate": 62.5,
+                        "grad_norm": round(grad_norm, 2),
+                        "target_len_info": {"min": int(target_lengths.min().item()), "max": int(target_lengths.max().item()), "avg": round(float(target_lengths.float().mean().item()), 1)},
+                        "ctc_spectrogram": []
+                    }
+                }
+                telemetry_pool.submit(async_post_telemetry, active_url, telemetry_payload)
+
+        del dataset
+        del dataloader
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    try:
+        streamer_stop_event.set()
+    except Exception:
+        pass
+    print("\n[+] Training session finished cleanly!")
+    print("[+] Syncing final Checkpoint back to Local Studio...")
+    sd_copy = {k: v.cpu().clone() for k, v in base_model.state_dict().items()}
+    telemetry_pool.submit(async_save_checkpoint, resolve_current_tunnel_url(TUNNEL_URL), SESSION_TOKEN, global_step, chunk_idx, sd_copy, VOCAB_CHARS, best_loss, best_loss, 1).result()
+    print("[+] Final Checkpoint synced successfully!")
+
+
+# =========================================================================
+# 8. MAIN ENTRYPOINT
+# =========================================================================
+def main():
+    num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
+    if IS_TPU_DEVICE:
+        print("=" * 70)
+        print("  AsterOCR Cloud TPU Worker: Detected XLA Environment!")
+        print("  Spawning PyTorch XLA Multi-Process (8 Cores)...")
+        print("=" * 70)
+        try:
+            try:
+                import torch_xla.distributed.xla_multiprocessing as xmp
+            except ImportError:
+                import torch_xla.distributed.xmp as xmp
+            # Set environment variables for TPU
+            os.environ['XLA_USE_BF16'] = '1'
+            os.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '1000000000'
+            xmp.spawn(xla_ddp_worker, args=(), nprocs=None, start_method='fork')
+        except Exception as e:
+            import traceback
+            print(f"[!] Lỗi xmp.spawn: {e}")
+            traceback.print_exc()
+            print("[!] Fallback to Single Device...")
+            single_device_worker()
+    elif num_gpus >= 2:
+        print("=" * 70)
+        print(f"  AsterOCR Cloud GPU Worker: Detected {num_gpus}x GPUs!")
+        print(f"  Spawning PyTorch Multi-Process DistributedDataParallel (DDP)...")
+        print("=" * 70)
+        try:
+            mp.set_start_method("spawn", force=True)
+        except Exception:
+            pass
+        mp.spawn(ddp_train_worker, nprocs=num_gpus, args=(num_gpus,))
+    else:
+        single_device_worker()
+
+if __name__ == "__main__":
+    main()