Attempting new file structure that requires fewer drastic changes

model_index.json

@@ -1,20 +1,11 @@
 {
   "_class_name": "TextConditionalDDPMPipeline",
   "_diffusers_version": "0.32.2",
-  "scheduler": [
-    "diffusers",
-    "DDPMScheduler"
-  ],
-  "text_encoder": [
-    "models.text_model",
-    "TransformerModel"
-  ],
-  "tokenizer": [
-    "tokenizer",
-    "Tokenizer"
-  ],
-  "unet": [
-    "diffusers",
-    "UNet2DConditionModel"
-  ]
-}
+  "custom_pipeline": "models/text_diffusion_pipeline.py",
+  "components": {
+    "unet": { "type": "UNet2DConditionModel", "subfolder": "unet" },
+    "text_encoder": { "type": "models.text_model.TransformerModel", "subfolder": "text_encoder" },
+    "tokenizer": { "type": "Tokenizer", "file": "tokenizer.py" },
+    "scheduler": { "type": "DDPMScheduler", "subfolder": "scheduler" }
+  }
+}

models/text_diffusion_pipeline.py

@@ -0,0 +1,442 @@
+import torch
+import torch.nn.functional as F
+from typing import NamedTuple, Optional
+import os
+from diffusers import DDPMPipeline, UNet2DConditionModel, DDPMScheduler
+import json
+# Running the main at the end of this requires messing with this import
+from models.text_model import TransformerModel  
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+import util.common_settings as common_settings
+import models.sentence_transformers_helper as st_helper
+import models.text_model as text_model
+#from models.general_training_helper import get_scene_from_embeddings
+            
+class PipelineOutput(NamedTuple):
+    images: torch.Tensor
+    
+
+
+# Create a custom pipeline for text-conditional generation
+class TextConditionalDDPMPipeline(DDPMPipeline):
+    def __init__(self, unet, scheduler, text_encoder=None, tokenizer=None, supports_pretrained_split=False, block_embeddings=None):
+        super().__init__(unet=unet, scheduler=scheduler)
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.supports_negative_prompt = hasattr(unet, 'negative_prompt_support') and unet.negative_prompt_support
+        self.supports_pretrained_split = supports_pretrained_split
+        self.block_embeddings = block_embeddings
+
+        if self.tokenizer is None and self.text_encoder is not None:
+            # Use the tokenizer from the text encoder if not provided
+            self.tokenizer = self.text_encoder.tokenizer
+        
+        # Register the text_encoder so that .to(), .cpu(), .cuda(), etc. work correctly
+        self.register_modules(
+            unet=unet,
+            scheduler=scheduler,
+            text_encoder=self.text_encoder,
+            tokenizer=self.tokenizer,
+        )
+    
+    # Override the to() method to ensure text_encoder is moved to the correct device
+    def to(self, device=None, dtype=None):
+        # Call the parent's to() method first
+        pipeline = super().to(device, dtype)
+        
+        # Additionally move the text_encoder to the device
+        if self.text_encoder is not None:
+            self.text_encoder.to(device)
+        
+        return pipeline
+
+    def save_pretrained(self, save_directory):
+        os.makedirs(save_directory, exist_ok=True)
+        super().save_pretrained(save_directory)  # saves UNet and scheduler
+
+        # Save block_embeddings tensor if it exists
+        if self.block_embeddings is not None:
+            torch.save(self.block_embeddings, os.path.join(save_directory, "block_embeddings.pt"))
+
+        # Save supports_negative_prompt and supports_pretrained_split flags
+        with open(os.path.join(save_directory, "pipeline_config.json"), "w") as f:
+            json.dump({
+                "supports_negative_prompt": self.supports_negative_prompt,
+                "supports_pretrained_split": self.supports_pretrained_split,
+                "text_encoder_type": type(self.text_encoder).__name__   
+            }, f)
+
+
+        #Text encoder/tokenizer saving is different depending on if we're using a larger pretrained model
+        if isinstance(self.text_encoder, TransformerModel):
+            # Save custom text encoder
+            if self.text_encoder is not None:
+                self.text_encoder.save_pretrained(os.path.join(save_directory, "text_encoder"))
+        else:
+            #Save pretrained tokenizer by name, so we can load from huggingface instead of saving a giant local model
+            text_encoder_info = {
+                "text_encoder_name": self.text_encoder.config.name_or_path,
+                "tokenizer_name": self.tokenizer.name_or_path,
+            }
+
+            text_encoder_directory = os.path.join(save_directory, "text_encoder")
+            os.makedirs(text_encoder_directory, exist_ok=True)
+
+            with open(os.path.join(text_encoder_directory, "loading_info.json"), "w") as f:
+                json.dump(text_encoder_info, f)
+            
+
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, **kwargs):
+        #from diffusers.utils import load_config, load_state_dict
+        # Load model_index.json
+        #model_index = load_config(pretrained_model_path)
+
+        # Load components manually
+        unet_path = os.path.join(pretrained_model_path, "unet")
+        unet = UNet2DConditionModel.from_pretrained(unet_path)
+
+        scheduler_path = os.path.join(pretrained_model_path, "scheduler")
+        # Have heard that DDIMScheduler might be faster for inference, though not necessarily better
+        scheduler = DDPMScheduler.from_pretrained(scheduler_path)
+
+        tokenizer = None
+        text_encoder_path = os.path.join(pretrained_model_path, "text_encoder")
+        
+        if os.path.exists(text_encoder_path):
+            #Test for the new saving system, where we save a simple config file
+            if os.path.exists(os.path.join(text_encoder_path, "loading_info.json")):
+                with open(os.path.join(text_encoder_path, "loading_info.json"), "r") as f:
+                    encoder_config = json.load(f)
+
+                text_encoder = AutoModel.from_pretrained(encoder_config['text_encoder_name'], trust_remote_code=True)
+                tokenizer = AutoTokenizer.from_pretrained(encoder_config['tokenizer_name'])
+            
+            #Legacy loading system, loads models directly if the whole thing is saved in the directory
+            else:
+                try:
+                    text_encoder = AutoModel.from_pretrained(text_encoder_path, local_files_only=True, trust_remote_code=True)
+                    tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, local_files_only=True)
+                except (ValueError, KeyError):
+                    text_encoder = TransformerModel.from_pretrained(text_encoder_path)
+                    tokenizer = text_encoder.tokenizer
+        else:
+            text_encoder = None
+
+        # Instantiate your pipeline
+        pipeline = cls(
+            unet=unet,
+            scheduler=scheduler,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            **kwargs,
+        )
+
+        #Loads block embeddings if present
+        block_embeds_path = os.path.join(pretrained_model_path, "block_embeddings.pt")
+        if os.path.exists(block_embeds_path):
+            pipeline.block_embeddings = torch.load(block_embeds_path, map_location="cpu")
+        else:
+            pipeline.block_embeddings = None
+        
+
+        # Load supports_negative_prompt flag if present
+        config_path = os.path.join(pretrained_model_path, "pipeline_config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                config = json.load(f)
+            pipeline.supports_negative_prompt = config.get("supports_negative_prompt", False)
+            pipeline.supports_pretrained_split = config.get("supports_pretrained_split", False)
+        return pipeline
+
+    # --- Handle batching for captions ---
+    def _prepare_text_batch(self, text: Optional[str | list[str]], batch_size: int, name: str) -> Optional[list[str]]:
+        if text is None:
+            return None
+        if isinstance(text, str):
+            return [text] * batch_size
+        if isinstance(text, list):
+            if len(text) == 1:
+                return text * batch_size
+            if len(text) != batch_size:
+                raise ValueError(f"{name} list length {len(text)} does not match batch_size {batch_size}")
+            return text
+        raise ValueError(f"{name} must be a string or list of strings")
+
+    def _prepare_initial_sample(self, 
+                                raw_latent_sample: Optional[torch.Tensor],
+                                input_scene: Optional[torch.Tensor],
+                                batch_size: int, height: int, width: int,
+                                generator: Optional[torch.Generator]) -> torch.Tensor:
+        """Prepare the initial sample for diffusion."""
+ 
+        sample_shape = (batch_size, self.unet.config.in_channels, height, width)
+
+        if raw_latent_sample is not None:
+            if input_scene is not None:
+                raise ValueError("Cannot provide both raw_latent_sample and input_scene")
+            sample = raw_latent_sample.to(self.device)
+            if sample.shape[1] != sample_shape[1]:
+                raise ValueError(f"Wrong number of channels in raw_latent_sample: Expected {self.unet.config.in_channels} but got {sample.shape[1]}")
+            if sample.shape[0] == 1 and batch_size > 1:
+                sample = sample.repeat(batch_size, 1, 1, 1)
+            elif sample.shape[0] != batch_size:
+                raise ValueError(f"raw_latent_sample batch size {sample.shape[0]} does not match batch_size {batch_size}")
+        elif input_scene is not None:
+            # input_scene can be (H, W) or (batch_size, H, W)
+            scene_tensor = torch.tensor(input_scene, dtype=torch.long, device=self.device)
+            if scene_tensor.dim() == 2:
+                # (H, W) -> repeat for batch
+                scene_tensor = scene_tensor.unsqueeze(0).repeat(batch_size, 1, 1)
+            elif scene_tensor.shape[0] == 1 and batch_size > 1:
+                scene_tensor = scene_tensor.repeat(batch_size, 1, 1)
+            elif scene_tensor.shape[0] != batch_size:
+                raise ValueError(f"input_scene batch size {scene_tensor.shape[0]} does not match batch_size {batch_size}")
+            # One-hot encode: (batch, H, W, C)
+            one_hot = F.one_hot(scene_tensor, num_classes=self.unet.config.in_channels).float()
+            # (batch, H, W, C) -> (batch, C, H, W)
+            sample = one_hot.permute(0, 3, 1, 2)
+        else:
+            # Start from random noise
+            sample = torch.randn(sample_shape, generator=generator, device=self.device)
+
+        return sample
+
+    def __call__(
+        self,
+        caption: Optional[str | list[str]] = None,
+        negative_prompt: Optional[str | list[str]] = None,
+        generator: Optional[torch.Generator] = None,
+        num_inference_steps: int = common_settings.NUM_INFERENCE_STEPS,
+        guidance_scale: float = common_settings.GUIDANCE_SCALE,
+        height: int = common_settings.MARIO_HEIGHT,
+        width: int = common_settings.MARIO_WIDTH,
+        raw_latent_sample: Optional[torch.FloatTensor] = None,
+        input_scene: Optional[torch.Tensor] = None,
+        output_type: str = "tensor",
+        batch_size: int = 1,
+        show_progress_bar: bool = True,
+    ) -> PipelineOutput:
+        """Generate a batch of images based on text input using the diffusion model.
+
+        Args:
+            caption: Text description(s) of the desired output. Can be a string or list of strings.
+            negative_prompt: Text description(s) of what should not appear in the output. String or list.
+            generator: Random number generator for reproducibility.
+            num_inference_steps: Number of denoising steps (more = higher quality, slower).
+            guidance_scale: How strongly the generation follows the text prompt (higher = stronger).
+            height: Height of generated image in tiles.
+            width: Width of generated image in tiles.
+            raw_latent_sample: Optional starting point for diffusion instead of random noise.
+                Must have correct number of channels matching the UNet.
+            input_scene: Optional 2D or 3D int tensor where each value corresponds to a tile type.
+                Will be converted to one-hot encoding as starting point.
+            output_type: Currently only "tensor" is supported.
+            batch_size: Number of samples to generate in parallel.
+
+        Returns:
+            PipelineOutput containing the generated image tensor (batch_size, ...).
+        """
+
+        #       I would like to simplify the code to this, but the AI suggestion didn't work, and 
+        #       I did not feel good just pasting it all in. Will need to tackle it bit by bit.
+
+        #        if caption is not None and self.text_encoder is None:
+        #            raise ValueError("Text encoder required for conditional generation")
+    
+        #        self.unet.eval()
+        #        if self.text_encoder is not None:
+        #            self.text_encoder.to(self.device)
+        #            self.text_encoder.eval()
+        #
+        #        with torch.no_grad():
+        #            # Process text inputs
+        #            captions = self.prepare_text_batch(caption, batch_size, "caption")
+        #            negatives = self.prepare_text_batch(negative_prompt, batch_size, "negative_prompt")
+         
+        #            # Get embeddings
+        #            text_embeddings = self.prepare_embeddings(captions, negatives, batch_size)
+        #            
+        #            # Set up initial latent state
+        #            sample = self.prepare_initial_sample(raw_latent_sample, input_scene, 
+        #                                              batch_size, height, width, generator)
+           
+        #            # Run diffusion process
+        #            sample = self.run_diffusion(sample, text_embeddings, num_inference_steps,
+        #                                      guidance_scale, generator, show_progress_bar,
+        #                                      has_caption=caption is not None,
+        #                                      has_negative=negative_prompt is not None)
+         
+        #            # Format output
+        #            if output_type == "tensor":
+        #                sample = F.softmax(sample, dim=1)
+        #            else:
+        #                raise ValueError(f"Unsupported output type: {output_type}")
+
+        #        return PipelineOutput(images=sample)
+
+        # Validate text encoder if we need it
+        if caption is not None and self.text_encoder is None:
+            raise ValueError("Text encoder is required for conditional generation")
+
+        self.unet.eval()
+        if self.text_encoder is not None:
+            self.text_encoder.to(self.device)
+            self.text_encoder.eval()
+
+        with torch.no_grad():
+            captions = self._prepare_text_batch(caption, batch_size, "caption")
+            negatives = self._prepare_text_batch(negative_prompt, batch_size, "negative_prompt")
+
+            # --- Prepare text embeddings ---
+            if(isinstance(self.text_encoder, TransformerModel)):
+                text_embeddings = text_model.get_embeddings(batch_size=batch_size,
+                                                            tokenizer=self.text_encoder.tokenizer,
+                                                            text_encoder=self.text_encoder,
+                                                            captions=captions,
+                                                            neg_captions=negatives,
+                                                            device=self.device)
+            else: #Case for the pre-trained text encoder
+                if(self.supports_pretrained_split): #If we have a split flag incorporated
+                    text_embeddings = st_helper.get_embeddings_split(batch_size = batch_size,
+                                                            tokenizer=self.tokenizer,
+                                                            model=self.text_encoder,
+                                                            captions=captions,
+                                                            neg_captions=negatives,
+                                                            device=self.device)
+                else:
+                    text_embeddings = st_helper.get_embeddings(batch_size = batch_size,
+                                                                tokenizer=self.tokenizer,
+                                                                model=self.text_encoder,
+                                                                captions=captions,
+                                                                neg_captions=negatives,
+                                                                device=self.device)
+
+                            
+            # --- Set up initial latent state ---
+            sample = self._prepare_initial_sample(raw_latent_sample, input_scene, 
+                                                 batch_size, height, width, generator)
+
+            # --- Set up diffusion process ---
+            self.scheduler.set_timesteps(num_inference_steps)
+
+            # Denoising loop
+            iterator = self.progress_bar(self.scheduler.timesteps) if show_progress_bar else self.scheduler.timesteps
+            for t in iterator:
+                # Handle conditional generation
+                if captions is not None:
+                    if negatives is not None:
+                        # Three copies for negative prompt guidance
+                        model_input = torch.cat([sample, sample, sample], dim=0)
+                    else:
+                        # Two copies for standard classifier-free guidance
+                        model_input = torch.cat([sample, sample], dim=0)
+                else:
+                    model_input = sample
+
+                # Predict noise residual
+                model_kwargs = {"encoder_hidden_states": text_embeddings}
+                noise_pred = self.unet(model_input, t, **model_kwargs).sample
+
+                # Apply guidance
+                if captions is not None:
+                    if negatives is not None:
+                        # Split predictions for negative, unconditional, and text-conditional
+                        noise_pred_neg, noise_pred_uncond, noise_pred_text = noise_pred.chunk(3)
+                        noise_pred_guided = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                        noise_pred = noise_pred_guided - guidance_scale * (noise_pred_neg - noise_pred_uncond)
+                    else:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # Compute previous sample: x_{t-1} = scheduler(x_t, noise_pred)
+                sample = self.scheduler.step(noise_pred, t, sample, generator=generator).prev_sample
+
+            # Convert to output format
+            if output_type == "tensor":
+                if self.block_embeddings is not None:
+                    sample = get_scene_from_embeddings(sample, self.block_embeddings)
+                else:
+                    # Apply softmax to get probabilities for each tile type
+                    sample = F.softmax(sample, dim=1)
+                    sample = sample.detach().cpu() 
+            else:
+                raise ValueError(f"Unsupported output type: {output_type}")
+
+        return PipelineOutput(images=sample)
+
+    def print_unet_architecture(self):
+        """Prints the architecture of the UNet model."""
+        print(self.unet)
+
+    def print_text_encoder_architecture(self):
+        """Prints the architecture of the text encoder model, if it exists."""
+        if self.text_encoder is not None:
+            print(self.text_encoder)
+        else:
+            print("No text encoder is set.")
+
+    def save_unet_architecture_pdf(self, height, width, filename="unet_architecture", batch_size=1, device=None):
+        """
+        Have to separately install torchview for this to work
+
+        Saves a visualization of the UNet architecture as a PDF using torchview.
+        Args:
+            height: Height of the dummy input.
+            width: Width of the dummy input.
+            filename: Output PDF filename.
+            batch_size: Batch size for dummy input.
+            device: Device to run the dummy input on (defaults to pipeline device).
+        """
+        from torchview import draw_graph
+        import graphviz
+        
+        if device is None:
+            device = self.device if hasattr(self, 'device') else 'cpu'
+        in_channels = self.unet.config.in_channels if hasattr(self.unet, 'config') else 1
+        sample_shape = tuple([batch_size, in_channels, height, width])
+
+        dummy_x = torch.randn(size=sample_shape, device=device)
+        dummy_t = torch.tensor([0] * batch_size, dtype=torch.long, device=device)
+
+        # Prepare dummy text embedding (match what your UNet expects)
+        if hasattr(self.unet, 'config') and hasattr(self.unet.config, 'cross_attention_dim'):
+            cross_attention_dim = self.unet.config.cross_attention_dim
+        else:
+            cross_attention_dim = 128  # fallback
+        encoder_hidden_states = torch.randn(batch_size, 1, cross_attention_dim, device=device)
+
+        self.unet.eval()
+        inputs = (dummy_x, dummy_t, encoder_hidden_states)
+        #self.unet.down_blocks = self.unet.down_blocks[:2]
+
+        graph = draw_graph(
+            model=self.unet,
+            input_data=inputs,
+            expand_nested=False,
+            #enable_output_shape=True,   
+            #roll_out="nested",
+            depth=1
+        )
+        #graph.visual_graph.engine = "neato"
+        graph.visual_graph.attr(#rankdir="LR",
+                                nodesep="0.1",      # decrease space between nodes in the same rank (default ~0.25)
+                                ranksep="0.2",       # decrease space between ranks (default ~0.5)
+                                concentrate="true"  # merge edges between nodes in the same rank
+                            )
+        graph.visual_graph.node_attr.update(
+            shape="rectangle",
+            width="1.5",   # narrow width
+            height="0.5"  # taller height to make vertical rectangles
+            #fixedsize="true"
+        )
+        
+        graph.visual_graph.render(filename, format='pdf', cleanup=False)  # Cleanup removes intermediate files
+        graph.visual_graph.save('unet_architecture.dot')
+
+        # Save the graph to a PDF file
+        print(f"UNet architecture saved to {filename}")
+

models/text_model.py

@@ -0,0 +1,206 @@
+import argparse
+from xml.parsers.expat import model
+import torch
+import torch.nn as nn
+import math
+import os
+import json
+from safetensors.torch import save_file, load_file
+from tokenizer import Tokenizer
+
+def get_embeddings(batch_size, tokenizer, text_encoder, captions=None, neg_captions=None, device='cpu'):
+    max_length = text_encoder.max_seq_length
+    empty_ids = encode_token_captions([""] * batch_size, tokenizer, max_length, device=device)
+    embeddings = text_encoder.get_embeddings(empty_ids)
+
+    if(captions is not None):
+        caption_ids = encode_token_captions(captions, tokenizer, max_length, device=device)
+        caption_embeddings = text_encoder.get_embeddings(caption_ids)
+        embeddings = torch.cat((embeddings, caption_embeddings), dim=0)
+    
+    if(neg_captions is not None):
+        neg_ids = encode_token_captions(neg_captions, tokenizer, max_length, device=device)
+        neg_embeddings = text_encoder.get_embeddings(neg_ids)
+        embeddings = torch.cat((neg_embeddings, embeddings), dim=0)
+    
+    return embeddings.to(device)
+
+def encode_token_captions(captions, tokenizer, max_length, device='cpu'):
+    caption_ids = []
+    for caption in captions:
+        tokens = tokenizer.encode(caption)
+        caption_tokens = tokenizer.pad_sequence(tokens, max_length)
+        caption_ids.append(torch.tensor(caption_tokens, dtype=torch.long).unsqueeze(0))
+    return torch.cat(caption_ids, dim=0).to(device)
+
+
+
+
+
+
+
+
+
+# Transformer model for MLM training
+
+class TransformerModel(nn.Module):
+    def __init__(self, vocab_size, embedding_dim, hidden_dim, tokenizer=None, num_heads=8, num_layers=4, max_seq_length=100):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.max_seq_length = max_seq_length
+
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.positional_encoding = self.create_positional_encoding(max_seq_length, embedding_dim)
+
+        encoder_layers = nn.TransformerEncoderLayer(
+            d_model=embedding_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim,
+            batch_first=True
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layers, num_layers)
+        self.fc = nn.Linear(embedding_dim, vocab_size)
+
+        self.tokenizer = tokenizer
+
+    def create_positional_encoding(self, max_seq_length, embedding_dim):
+        # The implementation uses a sinusoidal positional encoding, which creates a unique pattern for each position in the sequence.
+        # The frequencies create unique values, the sin/cos bounds values
+        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
+        # Creates a set of divisors that create different frequencies
+        div_term = torch.exp(torch.arange(0, embedding_dim, 2).float() * (-math.log(10000.0) / embedding_dim))
+        pe = torch.zeros(max_seq_length, embedding_dim)
+        # Even dimensions use sin, odd dimensions use cos
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe.unsqueeze(0)
+
+    def get_embeddings(self, x):
+        """ This gets the actual latent embedding vectors """
+        # Ensure positional encoding is on the same device as input
+        pe = self.positional_encoding[:, :x.size(1), :].to(x.device)
+        # Embed input and add positional encoding
+        embedded = self.embedding(x) + pe
+        return self.transformer(embedded)
+
+    def forward(self, x):
+        """ This gets the token within the vocabulary """
+        transformer_out = self.get_embeddings(x)
+        # Project to vocabulary size
+        return self.fc(transformer_out)
+
+    def save_pretrained(self, save_directory):
+        os.makedirs(save_directory, exist_ok=True)
+
+        config = {
+            "vocab_size": self.vocab_size,
+            "embedding_dim": self.embedding_dim,
+            "hidden_dim": self.hidden_dim,
+            "num_heads": self.num_heads,
+            "num_layers": self.num_layers,
+            "max_seq_length": self.max_seq_length,
+        }
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
+            json.dump(config, f)
+
+        # Save model weights
+        save_file(self.state_dict(), os.path.join(save_directory, "model.safetensors"))
+
+        # Save tokenizer if present
+        if self.tokenizer is not None:
+            self.tokenizer.save(os.path.join(save_directory, "tokenizer.pkl"))
+
+    @classmethod
+    def from_pretrained(cls, load_directory):
+        with open(os.path.join(load_directory, "config.json")) as f:
+            config = json.load(f)
+
+        model = cls(**config)
+
+        # Load weights
+        state_dict = load_file(os.path.join(load_directory, "model.safetensors"))
+        model.load_state_dict(state_dict)
+
+        # Load tokenizer if available
+        tokenizer_path = os.path.join(load_directory, "tokenizer.pkl")
+        if os.path.exists(tokenizer_path):
+            tokenizer = Tokenizer()
+            tokenizer.load(tokenizer_path)
+            model.tokenizer = tokenizer
+
+        return model
+    
+    def print_architecture(self, inputs=None):
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--model_path", type=str, required=True, help="Path to trained transformer model")
+        parser.add_argument("--json", type=str, default="SMB1_LevelsAndCaptions-regular-test.json", help="Path to dataset json file")
+        parser.add_argument("--num_samples", type=int, default=10, help="Number of captions to evaluate")
+        parser.add_argument("--mask_prob", type=float, default=0.15, help="Probability of masking each token")
+
+        parser.add_argument("--compare_checkpoints", action="store_true", default=False, help="Run comparison across all model checkpoints")
+        args = parser.parse_args()
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = TransformerModel.from_pretrained(args.model_path).to(device)
+        print(f"Loaded model from {args.model_path}")
+
+        import os
+        import re
+        import json
+        import matplotlib.pyplot as plt
+        from torchview import draw_graph
+        import graphviz
+
+        graph = draw_graph(
+            model=model,
+            input_data=inputs,
+            expand_nested=False,
+            #enable_output_shape=True,   
+            #roll_out="nested",
+            depth=1
+        )
+
+        # Save plot
+        filename = 'mlm_architecture'
+        graph.visual_graph.render(filename, format='pdf', cleanup=False)  # Cleanup removes intermediate files
+        #graph.visual_graph.save('unet_architecture.dot')
+
+    def save_architecture_pdf(self, filename="transformer_architecture.pdf", input_length=32):
+        """Save a visualization of the model architecture as a PDF using torchview."""
+        try:
+            from torchview import draw_graph
+        except ImportError:
+            raise ImportError("torchview is required for model visualization. Install with 'pip install torchview'.")
+        import torch
+        import os
+        # Create a dummy input of the correct type for the model
+        captions = ["full floor. two coins. one pipe.", "floor with two gaps. one cannon. many enemies."]
+        tensor = encode_token_captions(captions, self.tokenizer, self.max_seq_length, device=next(self.parameters()).device)
+        input_length = tensor.size(1) if tensor.dim() > 1 else self.max_seq_length
+
+        num_tokens_list = [len(self.tokenizer.encode(c)) for c in captions]
+        input_length = max(num_tokens_list) if num_tokens_list else input_length
+        dummy_input = torch.zeros((1, input_length), dtype=torch.long, device=next(self.parameters()).device)
+
+        # Draw the graph and save as PNG
+        graph = draw_graph(self, input_data=dummy_input, expand_nested=True, save_graph=True, filename=filename.replace('.pdf',''), directory=".", depth=2)
+        png_file = filename.replace('.pdf', '.png')
+        # Convert PNG to PDF
+        if os.path.exists(png_file):
+            try:
+                from PIL import Image
+                im = Image.open(png_file)
+                im.save(filename, "PDF", resolution=100.0)
+                print(f"Saved architecture PDF to {filename}")
+                # Optionally, remove the PNG file
+                os.remove(png_file)
+            except ImportError:
+                print(f"PIL not installed. Architecture saved as PNG: {png_file}")
+            except Exception as e:
+                print(f"Could not convert PNG to PDF: {e}")
+        else:
+            print(f"Could not find PNG file to convert: {png_file}")

tokenizer.py

@@ -0,0 +1,147 @@
+import json
+import re
+from collections import Counter
+import pickle
+import argparse
+
+class Tokenizer:
+    def __init__(self):
+        self.special_tokens = ["[PAD]", "[MASK]"]
+        self.vocab = {}
+        self.token_to_id = {}
+        self.id_to_token = {}
+
+    def tokenize(self, text):
+        # Match words, numbers, periods, and commas as separate tokens
+        tokens = re.findall(r'\w+|[.,]|\[mask\]|\[pad\]', text.lower())
+        # Restore MASK and PAD to all caps
+        modified_list = []
+        for s in tokens:
+            modified_s = s.replace("[mask]", "[MASK]").replace("[pad]", "[PAD]")
+            modified_list.append(modified_s)
+        return modified_list
+
+    def pad_sequence(self, tokens, length):
+        """Pads tokenized sequences to length with a padding token (assumed to be '[PAD]')."""
+        if len(tokens) > length:
+            raise ValueError(f"Token sequence length {len(tokens)} exceeds specified length {length}.")
+        
+        pad_token = self.token_to_id["[PAD]"]
+        return tokens + [pad_token] * (length - len(tokens))
+
+    def build_vocab(self, dataset_path, min_freq=1):
+        token_counter = Counter()
+
+        with open(dataset_path, 'r') as f:
+            data = json.load(f)
+            for entry in data:
+                caption = entry['caption']
+                tokens = self.tokenize(caption)
+                token_counter.update(tokens)
+
+        # Keep tokens that meet the min frequency
+        tokens = [tok for tok, count in token_counter.items() if count >= min_freq]
+
+        # Ensure special tokens are always included
+        all_tokens = self.special_tokens + sorted(tokens)
+        
+        # Build vocab dictionaries
+        self.vocab = {tok: idx for idx, tok in enumerate(all_tokens)}
+        self.token_to_id = self.vocab
+        self.id_to_token = {idx: tok for tok, idx in self.vocab.items()}
+
+        print(f"Vocabulary size: {len(self.vocab)}")
+
+    def encode(self, text):
+        tokens = self.tokenize(text)
+        encoded = []
+        for tok in tokens:
+            if tok not in self.token_to_id:
+                raise ValueError(f"Unknown token encountered: {tok} in {text}")
+            encoded.append(self.token_to_id[tok])
+        return encoded
+
+    def encode_batch(self, texts, pad_to_length=None):
+        """
+        Encode a batch of texts into token IDs with padding to ensure uniform length.
+    
+        Args:
+            texts (list): A list of strings to encode
+            pad_to_length (int, optional): Length to pad all sequences to. If None,
+                                          will pad to the length of the longest sequence.
+        
+        Returns:
+            list: A list of lists, where each inner list contains the token IDs for a text
+        """
+        # Get the padding token ID
+        pad_token = self.token_to_id["[PAD]"]
+    
+        # First encode all texts
+        encoded_texts = []
+        for text in texts:
+            try:
+                encoded = self.encode(text)
+                encoded_texts.append(encoded)
+            except ValueError as e:
+                raise ValueError(f"Error encoding text: {text}. {str(e)}")
+    
+        # Determine padding length
+        if pad_to_length is None:
+            pad_to_length = max(len(seq) for seq in encoded_texts)
+    
+        # Pad sequences to uniform length
+        padded_texts = []
+        for seq in encoded_texts:
+            if len(seq) > pad_to_length:
+                # Truncate if too long
+                padded_texts.append(seq[:pad_to_length])
+            else:
+                # Pad if too short
+                padding = [pad_token] * (pad_to_length - len(seq))
+                padded_texts.append(seq + padding)
+    
+        return padded_texts
+
+    def decode(self, token_ids):
+        return ' '.join(self.id_to_token[tok_id] for tok_id in token_ids)
+
+    def save(self, path):
+        with open(path, 'wb') as f:
+            pickle.dump({'vocab': self.vocab}, f)
+
+    def load(self, path):
+        with open(path, 'rb') as f:
+            data = pickle.load(f)
+            self.vocab = data['vocab']
+            self.token_to_id = self.vocab
+            self.id_to_token = {idx: tok for tok, idx in self.vocab.items()}
+
+    def get_vocab(self):
+        return sorted(self.vocab.keys())
+
+    def get_vocab_size(self):
+        return len(self.vocab)
+
+if __name__ == "__main__":
+    tokenizer = Tokenizer()
+
+    parser = argparse.ArgumentParser(description="Tokenizer utility for saving and loading vocabularies.")
+    parser.add_argument("action", choices=["save", "load"], help="Action to perform: 'save' or 'load'.")
+    parser.add_argument("--json_file", type=str, default='Mario_LevelsAndCaptions.json', help="Path to the JSON file containing the dataset (required for 'save').")
+    parser.add_argument("--pkl_file", type=str, default='Mario_Tokenizer.pkl', help="Path to the pickle file to save/load the tokenizer.")
+
+    args = parser.parse_args()
+
+    if args.action == "save":
+        if not args.json_file:
+            raise ValueError("The --json_file argument is required for the 'save' action.")
+        tokenizer.build_vocab(args.json_file)
+        tokenizer.save(args.pkl_file)
+    elif args.action == "load":
+        tokenizer.load(args.pkl_file)
+
+    # Example usage
+    #print(tokenizer.encode("floor with one gap. one enemy."))
+    #print(tokenizer.get_vocab())
+    #for id, token in tokenizer.id_to_token.items():
+    #    print(id,":",token)

util/common_settings.py

@@ -0,0 +1,18 @@
+
+NUM_INFERENCE_STEPS = 30
+GUIDANCE_SCALE = 7.5
+
+MARIO_HEIGHT = 16
+MARIO_WIDTH = 16
+
+MARIO_TILE_PIXEL_DIM = 16
+MARIO_TILE_COUNT = 13
+
+LR_HEIGHT = 32
+LR_WIDTH = 32
+
+LR_TILE_PIXEL_DIM = 8
+LR_TILE_COUNT = 8
+
+MEGAMAN_HEIGHT = 14
+MEGAMAN_WIDTH = 16