Upload eruku_continuous_inf.py with huggingface_hub

eruku_continuous_inf.py

@@ -0,0 +1,569 @@
+import torch
+import os as _os
+from transformers import AutoTokenizer
+from transformers import T5ForConditionalGeneration, T5Config
+from custom_datasets import HFDataCollector
+from einops.layers.torch import Rearrange
+from einops import rearrange, repeat
+from torch.nn import MSELoss, CTCLoss, CrossEntropyLoss
+from pathlib import Path
+from torchvision.utils import make_grid, save_image
+from PIL import Image, ImageDraw, ImageFont
+from models.origami import OrigamiNet
+from diffusers import AutoencoderKL
+from torch.nn.utils.rnn import pad_sequence
+from torchvision.transforms import Normalize
+import numpy as np
+import torch.nn as nn
+from typing import Tuple
+
+# Safer defaults for clearer NCCL/CUDA error reporting during debugging
+_os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
+_os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1")
+_os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+
+
+def _safe_int_from_maybe_tensor(value, fallback_min: int = 64) -> int:
+    """Convert a python int or 0-dim tensor (cpu/cuda) to int safely.
+
+    - Synchronizes CUDA before .item() to surface the true failing kernel site
+    - Moves to CPU before scalarization
+    - Falls back to a reasonable minimum on unexpected errors
+    """
+    try:
+        if isinstance(value, torch.Tensor):
+            scalar_tensor = value
+            # Take the first element if tensor is not scalar
+            if scalar_tensor.dim() > 0:
+                scalar_tensor = scalar_tensor.reshape(-1)[0]
+            # Synchronize to attribute errors to the right op during debug
+            if scalar_tensor.is_cuda:
+                try:
+                    torch.cuda.synchronize(scalar_tensor.device)
+                except Exception:
+                    pass
+            return int(scalar_tensor.detach().to("cpu").item())
+        return int(value)
+    except Exception:
+        # As a last resort, return a conservative minimum width
+        return int(fallback_min)
+
+def pad_images(images, padding_value=1):
+    images = [rearrange(img, 'c h w -> w c h') for img in images]
+    padded = rearrange(pad_sequence(images, padding_value=padding_value), 'w b c h -> b c h w')
+    return padded.contiguous()
+
+
+
+
+# sog, eog, img
+SPECIAL_TOKEN_COUNT = 3
+
+class Emuru(torch.nn.Module):
+    def __init__(self, t5_checkpoint='google-t5/t5-base',
+                 vae_checkpoint='blowing-up-groundhogs/emuru_vae',
+                 ocr_checkpoint='files/checkpoints/Origami_bw_img/origami.pth', slices_per_query=1, channels=1, text_dropout_probability=0.0, img_dropout_probability=0.0):
+        super(Emuru, self).__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained('google/byt5-small')  # per-character tokenizer
+        self.tokenizer.add_tokens(["<sog>"])
+        self.data_collator = HFDataCollector(tokenizer=self.tokenizer)
+        self.t5_name_or_path = t5_checkpoint
+
+        self.padding_token = torch.tensor([[-0.4951,  0.8021,  0.3429,  0.5622,  0.5271,  0.5756,  0.7194,  0.6150]])
+        self.padding_token_threshold = 0.484982096850872
+
+        config = T5Config.from_pretrained(t5_checkpoint)
+        config.vocab_size = len(self.tokenizer)
+        self.T5 = T5ForConditionalGeneration(config)
+        # Expose a HF-like config for downstream trainers expecting model.config
+        self.config = self.T5.config
+        # Ensure a valid identifier is present for downstream AutoProcessor lookups
+        try:
+            if not getattr(self.config, "_name_or_path", None):
+                self.config._name_or_path = str(self.t5_name_or_path)
+        except Exception:
+            # As a safe fallback, set attribute directly
+            self.config._name_or_path = str(self.t5_name_or_path)
+        self.T5.lm_head = torch.nn.Identity()
+        self.normalize = Normalize(0.5, 0.5)
+        self.sos = torch.nn.Embedding(1, config.d_model)
+        self.sog = torch.nn.Embedding(1, config.d_model)
+        self.eog = torch.nn.Embedding(1, config.d_model)
+        
+        self.vae = AutoencoderKL.from_pretrained(vae_checkpoint)
+        
+        vae_latent_dim = 8 # self.vae.config.get('latent_channels', 8)
+
+        self.query_emb = torch.nn.Linear(vae_latent_dim * channels * slices_per_query, config.d_model)
+        self.t5_to_vae = torch.nn.Linear(config.d_model, vae_latent_dim * channels * slices_per_query)
+        self.t5_to_special = torch.nn.Linear(config.d_model, SPECIAL_TOKEN_COUNT)
+        self.t5_to_ocr = torch.nn.Linear(config.d_model, len(self.tokenizer), bias=False)
+
+        self.uncond_embedding = torch.nn.Embedding(1, config.d_model)
+        self.dropout_probability = 0.0
+        self.drop_text = False
+        self.drop_img = False
+
+        self.set_training(self.vae, False)
+
+        self.ocr = OrigamiNet.from_checkpoint(ocr_checkpoint, o_classes=165, n_channels=1)
+        self.set_training(self.ocr, False)
+        
+        self.query_rearrange = Rearrange('b c h (w q) -> b w (q c h)', q=slices_per_query)
+        self.special_rearrange = torch.nn.Identity()
+        # self.special_rearrange = Rearrange('b w (h c) -> b w (h c)')
+        self.z_rearrange = Rearrange('b w (q c h) -> b c h (w q)', c=channels, q=slices_per_query)
+        self.z_rearrange_eval = Rearrange('w b (q c h) -> b c h (w q)', c=channels, q=slices_per_query)
+
+        self.mse_criterion = MSELoss()#(reduction='none') # TODO:change reductions if you intend to add a mask
+        self.ce_criterion = CrossEntropyLoss()
+        # self.ctc_criterion = CTCLoss()
+        self.trainer = None
+        self.alpha = 1.0
+        # Minimal attributes for TRL compatibility
+        self.warnings_issued = {}
+        self._model_tags = set()
+
+    def add_model_tags(self, tags):
+        try:
+            if isinstance(tags, (list, tuple, set)):
+                self._model_tags.update(tags)
+            elif isinstance(tags, str):
+                self._model_tags.add(tags)
+        except Exception:
+            # No-op if tags updating fails
+            pass
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        """Enable gradient checkpointing - delegate to T5 model"""
+        if hasattr(self.T5, 'gradient_checkpointing_enable'):
+            self.T5.gradient_checkpointing_enable(gradient_checkpointing_kwargs)
+
+    def gradient_checkpointing_disable(self):
+        """Disable gradient checkpointing - delegate to T5 model"""
+        if hasattr(self.T5, 'gradient_checkpointing_disable'):
+            self.T5.gradient_checkpointing_disable()
+
+    def set_training(self, model, training):
+        model.train() if training else model.eval()
+        for param in model.parameters():
+            param.requires_grad = training 
+    
+    def _img_encode(self,img):
+        img = self.normalize(img)
+        # Ensure contiguous memory layout before encode to avoid kernel issues
+        img = img.contiguous()
+        return self.vae.encode(img.float()).latent_dist.sample()
+
+    @torch.no_grad()
+    def get_model_inputs(self, style_img, gen_img, style_len, gen_len, max_img_len):
+        bs = len(style_img)
+        decoder_inputs_embeds_list = []
+        specials_list = []
+        
+        # Move images to device and pad them
+        style_img = pad_images([el.to(self.T5.device) for el in style_img])
+        
+        if gen_img is not None:
+            gen_img = pad_images([el.to(self.T5.device) for el in gen_img])
+            gen_img_embeds = self._img_encode(gen_img)
+        else:
+            gen_img_embeds = None
+
+        style_img_embeds = self._img_encode(style_img)
+        
+        for el in range(bs):
+            if isinstance(style_len, int):
+                sl = style_len
+            else:
+                # Safely get scalar style length
+                sl_tensor = style_len[el] if hasattr(style_len, '__getitem__') else style_len
+                sl = _safe_int_from_maybe_tensor(sl_tensor)
+
+            # Ensure widths are within bounds
+            sl = max(64, min(sl, style_img_embeds.shape[-1]))
+            
+            # Start with style image embeds
+            sample_embeds_parts = [style_img_embeds[el,:,:,:sl//8]]
+            specials_parts = [torch.ones(sl//8) * 2] # Img token
+
+            if gen_img_embeds is not None and gen_len is not None:
+                if isinstance(gen_len, int):
+                    gl = gen_len
+                else:
+                    gl_tensor = gen_len[el] if hasattr(gen_len, '__getitem__') else gen_len
+                    gl = _safe_int_from_maybe_tensor(gl_tensor)
+
+                gl = max(64, min(gl, gen_img_embeds.shape[-1]))
+                sample_embeds_parts.extend([
+                    torch.ones(1, 8, 1).to(self.T5.device), # SOG token placeholder
+                    gen_img_embeds[el,:,:,:gl//8],
+                    torch.ones(1, 8, 1).to(self.T5.device), # EOG token placeholder
+                ])
+                specials_parts.extend([
+                    torch.zeros(1), # SOG
+                    torch.ones(gl//8) * 2, # Img
+                    torch.ones(1), # EOG
+                ])
+
+            sample_embeds = torch.cat(sample_embeds_parts, dim=-1)
+
+            h_dim = sample_embeds.shape[1]
+            sample_embeds = rearrange(sample_embeds, 'c h w -> w (h c)', h=h_dim, c=1)
+
+            decoder_inputs_embeds_list.append(sample_embeds)
+
+            sample_specials = torch.cat(specials_parts, dim=0).to(self.T5.device)
+            specials_list.append(sample_specials)
+
+        # Pad sequences and ensure consistent shapes
+        decoder_inputs_embeds_padded = pad_sequence(decoder_inputs_embeds_list, padding_value=1, batch_first=True)
+        specials_padded = pad_sequence(specials_list, padding_value=1, batch_first=True)
+        
+        # Ensure we don't exceed max_img_len
+        max_seq_len = max_img_len // 8
+        if decoder_inputs_embeds_padded.size(1) > max_seq_len:
+            decoder_inputs_embeds_padded = decoder_inputs_embeds_padded[:, :max_seq_len]
+        if specials_padded.size(1) > max_seq_len:
+            specials_padded = specials_padded[:, :max_seq_len]
+        
+        return {
+            'decoder_inputs_embeds': decoder_inputs_embeds_padded,
+            'specials': specials_padded.long(),
+        }
+
+    def forward(self, decoder_inputs_embeds_vae, specials, style_text, gen_text, ce_multiplier=1.0):
+        # style_img_embeds: [bs, w//8, 8, 1]
+        # generate text embeddings
+        
+        with torch.no_grad():
+            encoded_text = self.tokenizer([f"{style}<sog>{gen}" for style, gen in zip(style_text, gen_text)], padding=True, return_tensors="pt")
+        
+        # add special tokens to img
+        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=decoder_inputs_embeds_vae.size(0))
+        sog = repeat(self.sog.weight, '1 d -> b d', b=decoder_inputs_embeds_vae.size(0))
+        # eog = repeat(self.eog.weight, '1 d -> b d', b=decoder_inputs_embeds_vae.size(0))
+
+        decoder_inputs_embeds = self.query_emb(decoder_inputs_embeds_vae)
+        
+        # Fix the indexing assignment to avoid shape mismatch
+        sog_mask = (specials == 0)
+        eog_mask = (specials == 1)
+        
+        # Expand sog to match the sequence dimension
+        sog_expanded = sog.unsqueeze(1).expand(-1, decoder_inputs_embeds.size(1), -1)
+        
+        if sog_mask.any():
+            decoder_inputs_embeds[sog_mask] = sog_expanded[sog_mask]
+        
+        if eog_mask.any():
+            # Expand eog to match the sequence dimension
+            eog_expanded = self.eog.weight.unsqueeze(0).expand(decoder_inputs_embeds.size(0), decoder_inputs_embeds.size(1), -1)
+            decoder_inputs_embeds[eog_mask] = eog_expanded[eog_mask]
+
+        decoder_inputs_embeds = torch.cat(
+            [
+                sos,
+                decoder_inputs_embeds
+            ], dim = 1,
+        )
+
+        inputs_embeds = self.T5.shared(encoded_text['input_ids'].to(self.T5.device))
+        drop_ids = torch.rand(inputs_embeds.shape[0], device=inputs_embeds.device) < self.dropout_probability
+        if self.drop_text:
+            inputs_embeds = torch.where(drop_ids[:, None, None], self.uncond_embedding.weight, inputs_embeds)
+        if self.drop_img:
+            decoder_inputs_embeds = torch.where(drop_ids[:, None, None], self.uncond_embedding.weight, decoder_inputs_embeds)
+        
+        output = self.T5(inputs_embeds=inputs_embeds, attention_mask=encoded_text['attention_mask'].to(self.T5.device), decoder_inputs_embeds=decoder_inputs_embeds)
+        
+        vae_latent = self.t5_to_vae(output.logits[:, :-1])
+        special_latent = self.t5_to_special(output.logits[:, :-1]) # [bs, w//8, 3]
+        pred_latent = self.z_rearrange(vae_latent)
+        special_pred = self.special_rearrange(special_latent)
+
+        
+        ce_loss = ce_multiplier * self.ce_criterion(special_pred.flatten(0,1), specials.flatten(0,1))
+
+        mse_mask = (specials == 2).unsqueeze(2) # [bs, w//8] TODO:consider putting the mask back in
+        gt = decoder_inputs_embeds_vae * mse_mask
+        vae_latent = vae_latent * mse_mask
+        mse_loss = self.mse_criterion(vae_latent, gt)#/mse_mask.sum()
+        ocr_loss = 0
+
+        if self.alpha < 1.0:
+            pred_img = self.vae.decode(pred_latent).sample
+            gt_img = self.vae.decode(decoder_inputs_embeds_vae.unsqueeze(1)).sample
+            ocr_preds = self.ocr(pred_img)
+            ocr_gt = self.ocr(gt_img)
+            ocr_loss = self.mse_criterion(ocr_preds, ocr_gt)
+        else:
+            ocr_loss = torch.tensor(0.0).to(mse_loss.device)
+        loss = (ce_loss + mse_loss) * self.alpha + ocr_loss * (1 - self.alpha)
+        return {'loss': loss, 'mse_loss': mse_loss, 'ce_loss': ce_loss, 'ocr_loss': ocr_loss}, pred_latent
+
+    def split_characters(self, pred, gt, indices):
+        pred = self.vae.decode(pred).sample
+        gt = self.vae.decode(gt).sample
+        img = torch.cat([gt, pred], dim=-2)
+
+        curr_char = indices[0]
+        for idx, char in enumerate(indices):
+            if char != curr_char:
+                img[:, :, :, idx * 8 - 1] = -1
+                curr_char = char
+
+        img = self.write_text_below_image(img, self.tokenizer.decode(indices))
+
+        return img
+    
+
+    @torch.no_grad()
+    def write_text_below_image(self, image, text):
+        image = (torch.clamp(image, -1, 1) + 1) * 127.5
+        image = rearrange(image.to(torch.uint8), '1 1 h w -> h w').cpu().numpy()
+        image = Image.fromarray(image, mode='L')
+
+        text = text.replace('<pad>', '#').replace('</s>', '$')
+
+        # Load the font
+        font = ImageFont.load_default()
+        ascent, descent = font.getmetrics()
+        (width, baseline), (offset_x, offset_y) = font.font.getsize(text)
+
+        # Calculate dimensions for the new image
+        img_width, img_height = image.size
+        new_height = img_height + offset_y + ascent +descent
+
+        # Create a new image with white background
+        new_image = Image.new('L', (img_width, new_height), color='white')
+
+        # Paste the original image onto the new image
+        new_image.paste(image, (0, 0))
+
+        # Draw the text onto the new image
+        draw = ImageDraw.Draw(new_image)
+
+        curr_char = None
+        for idx, char in enumerate(text):
+            if char != curr_char:
+                curr_char = char
+                draw.text((idx * 8, img_height), char, fill='black', font=font)
+
+        return new_image
+    
+    @torch.inference_mode()
+    def generate(self, decoder_inputs_embeds_vae, style_text, gen_text, cfg_scale=1.0, max_new_tokens=64):
+        """
+        call this with bs=1 please
+        """
+        encoded_text = self.tokenizer([f"{style}<sog>{gen}" for style, gen in zip(style_text,gen_text)], padding=True, return_tensors="pt")
+        text_input_ids = encoded_text['input_ids'].to(self.T5.device)
+        text_mask = encoded_text['attention_mask'].to(self.T5.device)
+
+        sog = repeat(self.sog.weight, '1 d -> b 1 d', b=1)
+        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=1)
+        z_sequence = [decoder_inputs_embeds_vae]
+        special_sequence = torch.ones(decoder_inputs_embeds_vae.size(1))*3
+        if len(z_sequence) == 0:
+            decoder_inputs_embeds = sos
+        else:
+            decoder_inputs_embeds = self.query_emb(torch.cat(z_sequence, dim=1))
+            if len(style_text[0]) != 0:
+                decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
+            else:
+                decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds, sog], dim=1)
+                vae_latent = self.t5_to_vae(sog)
+                special_sequence = torch.cat([special_sequence, torch.zeros(1)])
+                z_sequence.append(vae_latent)
+
+        for i in range(max_new_tokens):
+            if cfg_scale != 1.0:
+                conditional_text_embeds = self.T5.shared(text_input_ids)
+                if self.drop_text:
+                    unconditional_text_embeds = torch.zeros_like(conditional_text_embeds).to(self.T5.device) + self.uncond_embedding.weight
+                else:
+                    unconditional_text_embeds = conditional_text_embeds
+
+                if self.drop_img:
+                    unconditional_decoder_inputs_embeds = torch.zeros_like(decoder_inputs_embeds).to(self.T5.device) + self.uncond_embedding.weight
+                else:
+                    unconditional_decoder_inputs_embeds = decoder_inputs_embeds
+
+                output_unconditional = self.T5(inputs_embeds=unconditional_text_embeds, attention_mask=text_mask, decoder_inputs_embeds=unconditional_decoder_inputs_embeds).logits[:, -1:]
+                output_conditional = self.T5(input_ids=text_input_ids, attention_mask=text_mask, decoder_inputs_embeds=decoder_inputs_embeds).logits[:, -1:]
+                output = output_unconditional + (output_conditional - output_unconditional) * cfg_scale
+            else:
+                output = self.T5(input_ids=text_input_ids, attention_mask=text_mask, decoder_inputs_embeds=decoder_inputs_embeds).logits[:, -1:]
+
+            special_prediction = self.t5_to_special(output)
+
+            if torch.argmax(special_prediction, dim=-1) == 0:
+                decoder_inputs_embeds = torch.cat([decoder_inputs_embeds, sog], dim=1)
+                vae_latent = self.t5_to_vae(output)
+                special_sequence = torch.cat([special_sequence, torch.zeros(1)])
+            elif torch.argmax(special_prediction, dim=-1) == 1:
+                special_sequence = torch.cat([special_sequence, torch.ones(1)])
+                vae_latent = self.t5_to_vae(output)
+                z_sequence.append(vae_latent)
+                break
+            else:
+                vae_latent = self.t5_to_vae(output)
+                decoder_inputs_embeds = torch.cat([decoder_inputs_embeds, self.query_emb(vae_latent)], dim=1)
+                special_sequence = torch.cat([special_sequence, torch.ones(1)*2])
+            z_sequence.append(vae_latent)
+            
+            
+        z_sequence = [el.to(self.vae.device) for el in z_sequence]
+        
+        z_sequence = torch.cat(z_sequence, dim=1)
+        img = torch.clamp(self.vae.decode(self.z_rearrange(z_sequence)).sample, -1, 1)
+        return img, special_sequence.to(self.T5.device)
+        
+    @torch.no_grad()
+    def continue_gen_test(self, gt, batch, max_new_tokens=64, cfg_scale=1.0):
+        gt = gt[:1]
+        def _continue_gen(style_len):
+            
+            generation = self.generate(batch['decoder_inputs_embeds'][:1, :style_len], batch['style_text'][:1], batch['gen_text'][:1], cfg_scale=cfg_scale, max_new_tokens=max_new_tokens)
+            test_img = generation[0]
+            special_sequence = generation[1].repeat_interleave(8)
+
+            
+            special_img = torch.zeros_like(test_img).repeat(1,3,1,1)
+            special_sequence = special_sequence[:special_img.size(-1)]
+            special_img[:,0,:,special_sequence == 2] = 1 # red: image
+            special_img[:,1,:,special_sequence == 0] = 1 # green: sog
+            special_img[:,2,:,special_sequence == 1] = 1 # blue: eog
+                        
+            try:
+                test_img[:, :, :, style_len * 8] = -1  # add a black line between style and pred
+            except:
+                print("couldn't add black line")
+                # add special_img to the bottom of test_img
+            test_img = torch.cat([test_img.repeat(1,3,1,1) , special_img], dim=-2)
+            return test_img
+        
+        gt = torch.clamp(self.vae.decode(gt).sample, -1, 1)
+        if type(batch['style_img_width']) == torch.Tensor:
+            style_img_width = batch['style_img_width'][0]
+        else:
+            style_img_width = batch['style_img_width']
+
+        return torch.cat(list(pad_images([
+            # make_grid(_continue_gen(style_img_width//8-10), nrow=1, normalize=True),
+            make_grid(_continue_gen(style_img_width//8), nrow=1, normalize=True),
+        ])), dim=-2)
+
+
+    def save_pretrained(self, path):
+        path = Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+        torch.save(self.T5.state_dict(), path / 'T5.pth')
+        torch.save(self.vae.state_dict(), path / 'VAE.pth')
+        torch.save(self.ocr.state_dict(), path / 'OCR.pth')
+        torch.save(self.query_emb.state_dict(), path / 'query_emb.pth')
+        torch.save(self.sos.state_dict(), path / 'sos.pth')
+
+    def load_pretrained(self, path):
+        path = Path(path)
+        self.T5.load_state_dict(torch.load(path / 'T5.pth'))
+        self.vae.load_state_dict(torch.load(path / 'VAE.pth'))
+        self.ocr.load_state_dict(torch.load(path / 'OCR.pth'))
+        self.query_emb.load_state_dict(torch.load(path / 'query_emb.pth'))
+        self.sos.load_state_dict(torch.load(path / 'sos.pth'))
+
+class DDPCompatibleEmuru(Emuru):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, batch_data, mode='train'):
+        """
+        Unified forward method that handles different modes for DDP compatibility
+        """
+        if mode == 'train':
+            # Training mode - expects the full batch with model inputs already computed
+            return super().forward(
+                batch_data['decoder_inputs_embeds'], 
+                batch_data['specials'], 
+                batch_data['style_text'], 
+                batch_data['gen_text']
+            )
+        elif mode == 'get_model_inputs':
+            # Mode to get model inputs
+            return super().get_model_inputs(
+                batch_data['style_img'],
+                batch_data['gen_img'], 
+                batch_data['style_img_width'],
+                batch_data['gen_img_width'],
+                batch_data['max_img_len']
+            )
+        elif mode == 'generate':
+            # Generation mode
+            return super().generate(
+                batch_data['decoder_inputs_embeds_vae'],
+                batch_data['style_text'],
+                batch_data['gen_text'],
+                batch_data.get('cfg_scale', 1.0),
+                batch_data.get('max_new_tokens', 64)
+            )
+        elif mode == 'continue_gen_test':
+            # Continue generation test mode
+            return super().continue_gen_test(
+                batch_data['gt'],
+                batch_data['batch'],
+                batch_data.get('cfg_scale', 1.0),
+                batch_data.get('max_new_tokens', 64)
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+    def module_get_model_inputs(self, style_img, gen_img, style_len, gen_len, max_img_len):
+        """Direct access method for get_model_inputs when not using DDP forward"""
+        return super().get_model_inputs(style_img, gen_img, style_len, gen_len, max_img_len)
+    
+    def module_continue_gen_test(self, gt, batch, max_new_tokens=64, cfg_scale=1.0):
+        """Direct access method for continue_gen_test when not using DDP forward"""
+        return super().continue_gen_test(gt, batch, max_new_tokens, cfg_scale)
+    
+    def module_vae_decode(self, latents):
+        """Direct access method for VAE decode"""
+        return self.vae.decode(latents)
+
+    def get_trainable_parameters(self):
+        """
+        Get only the parameters that have requires_grad=True
+        Useful for creating optimizers with only trainable parameters
+        """
+        return [p for p in self.parameters() if p.requires_grad]
+    
+    def get_parameter_count(self):
+        """
+        Get counts of total and trainable parameters
+        """
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {
+            'total_parameters': total_params,
+            'trainable_parameters': trainable_params,
+            'frozen_parameters': total_params - trainable_params
+        }
+    
+    def print_parameter_info(self):
+        """
+        Print detailed information about model parameters
+        """
+        info = self.get_parameter_count()
+        print(f"Model Parameter Info:")
+        print(f"  Total parameters: {info['total_parameters']:,}")
+        print(f"  Trainable parameters: {info['trainable_parameters']:,}")
+        print(f"  Frozen parameters: {info['frozen_parameters']:,}")
+        print(f"  Trainable ratio: {info['trainable_parameters']/info['total_parameters']:.2%}")
+        
+        # Print per-module info
+        print(f"\nPer-module breakdown:")
+        for name, module in self.named_children():
+            module_total = sum(p.numel() for p in module.parameters())
+            module_trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
+            if module_total > 0:
+                print(f"  {name}: {module_trainable:,}/{module_total:,} trainable ({module_trainable/module_total:.1%})")