Deleting directories, moving files into root

models/general_training_helper.py

@@ -1,172 +0,0 @@
-from torch.utils.data import DataLoader
-from level_dataset import LevelDataset
-import random
-from util.plotter import Plotter
-from datetime import datetime
-import os
-import threading
-import json
-import torch.nn.functional as F
-import torch
-
-
-
-
-def create_dataloaders(json_path, val_json, tokenizer, data_mode, augment, num_tiles, 
-                       negative_prompt_training, block_embeddings, batch_size):
-    # Initialize dataset
-    train_dataset = LevelDataset(
-        json_path=json_path,
-        tokenizer=tokenizer,
-        shuffle=True,
-        mode=data_mode,
-        augment=augment,
-        num_tiles=num_tiles,
-        negative_captions=negative_prompt_training,
-        block_embeddings=block_embeddings
-    )
-    val_dataset = None
-    if val_json is not None:
-        val_dataset = LevelDataset(
-            json_path=val_json,
-            tokenizer=tokenizer,
-            shuffle=False,
-            mode=data_mode,
-            augment=False,
-            num_tiles=num_tiles,
-            negative_captions=negative_prompt_training,
-            block_embeddings=block_embeddings
-        )
-
-    # Create dataloader
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_size=batch_size,
-        shuffle=True,
-        num_workers=4,
-        drop_last=True,
-        persistent_workers=True
-    )
-    
-    val_dataloader = None
-    if val_dataset is not None:
-        val_dataloader = DataLoader(
-            val_dataset,
-            batch_size=batch_size,
-            shuffle=False,
-            num_workers=4,
-            drop_last=False,
-            persistent_workers=True
-        )
-    
-    return train_dataloader, val_dataloader
-
-
-def get_random_training_samples(train_dataloader, negative_prompt_training, output_dir = None):
-    train_dataset = train_dataloader.dataset
-    # Sample four random captions from the dataset
-    sample_indices = [random.randint(0, len(train_dataset) - 1) for _ in range(4)]
-
-    sample_captions = [train_dataset[i][1] for i in sample_indices]
-    print("Sample captions:")
-    for caption in sample_captions:
-        print(caption)
-
-    sample_negative_captions = ""
-    if negative_prompt_training:
-        sample_negative_captions = [train_dataset[i][2] for i in sample_indices]
-        print("Sample negative captions:")
-        for caption in sample_negative_captions:
-            print(f"  NEG: {caption}")
-
-    #Write captions to a file
-    if output_dir is not None:
-        os.makedirs(output_dir, exist_ok=True)
-        out_path = os.path.join(output_dir, "sample_captions.txt")
-        with open(out_path, "w", encoding="utf-8") as f:
-            f.write("Sample captions:\n")
-            for caption in sample_captions:
-                f.write(str(caption) + "\n")
-            if negative_prompt_training:
-                f.write("\nSample negative captions:\n")
-                for caption in sample_negative_captions:
-                    f.write(str(caption) + "\n")
-        print(f"Sample captions written to {out_path}")
-
-
-    return sample_captions, sample_negative_captions
-
-
-def start_plotter(log_file, output_dir, left_key, right_key, left_label, right_label, png_name):
-    formatted_date = datetime.now().strftime(r'%Y%m%d-%H%M%S')
-
-    plotter = Plotter(log_file, update_interval=5.0, left_key=left_key, right_key=right_key,
-                            left_label=left_label, right_label=right_label, output_png=f'{png_name}_{formatted_date}.png')
-    plot_thread = threading.Thread(target=plotter.start_plotting)
-    plot_thread.daemon = True
-    plot_thread.start()
-    print(f"{png_name} plotting enabled. Progress will be saved to {os.path.join(output_dir, f'{png_name}_{formatted_date}.png')}")
-    return plotter, plot_thread
-
-
-def kill_plotter(plotter, plot_thread):
-    if plot_thread and plot_thread.is_alive():
-        plotter.stop_plotting()
-        plot_thread.join(timeout=5.0)
-        if plot_thread.is_alive():
-            print("Warning: Plot thread did not terminate properly")
-
-
-def load_config_from_json(config_path):
-    """Load hyperparameters from a JSON config file."""
-    try:
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-            print(f"Configuration loaded from {config_path}")
-            
-            # Print the loaded config for verification
-            print("Loaded hyperparameters:")
-            for key, value in config.items():
-                print(f"  {key}: {value}")
-                
-            return config
-    except (json.JSONDecodeError, FileNotFoundError) as e:
-        print(f"Error loading config file: {e}")
-        raise e
-
-
-def update_args_from_config(args, config):
-    """Update argparse namespace with values from config."""
-    # Convert config dict to argparse namespace
-    for key, value in config.items():
-        if hasattr(args, key):
-            setattr(args, key, value)
-    return args
-
-
-def get_scene_from_embeddings(image, block_embeddings):
-    """Code copied over from level_dataset, should give limited support for block embeddings"""
-    # Reshape sample to [batch_size * height * width, embedding_dim]
-    batch_size, embedding_dim, height, width = image.shape
-    
-    flat_samples = image.permute(0, 2, 3, 1).reshape(-1, embedding_dim)
-    
-    # Normalize vectors for cosine similarity
-    flat_samples = F.normalize(flat_samples, p=2, dim=1).cpu()
-    block_embeddings = F.normalize(block_embeddings, p=2, dim=1)
-
-    # Calculate cosine similarity between each position and all tile embeddings
-    similarities = torch.matmul(flat_samples, block_embeddings.t())
-    
-    # Get indices of most similar tiles
-    indices = torch.softmax(similarities, dim=1)
-    
-    
-    # Reshape back to [batch_size, height, width]
-    indices = indices.reshape(batch_size, height, width, 13)
-    indices = indices.permute(0, 3, 1, 2)
-
-    image=indices.detach().cpu()
-    return image
-
-

models/latent_diffusion_pipeline.py

@@ -1,99 +0,0 @@
-from diffusers import DDPMPipeline
-import torch
-import torch.nn.functional as F
-from typing import Optional, Union, List, Tuple
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.ddpm.pipeline_ddpm import ImagePipelineOutput
-import util.common_settings as common_settings
-import os
-import json
-from models.general_training_helper import get_scene_from_embeddings
-
-class UnconditionalDDPMPipeline(DDPMPipeline):
-    def __init__(self, unet, scheduler, block_embeddings=None):
-        super().__init__(unet, scheduler)
-
-        self.block_embeddings = block_embeddings
-    
-
-    def save_pretrained(self, save_directory):
-        os.makedirs(save_directory, exist_ok=True)
-        super().save_pretrained(save_directory)
-        # 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"))
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_path, **kwargs):
-        pipeline = super().from_pretrained(pretrained_model_path, **kwargs)
-        # Load block_embeddings tensor if it exists
-        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
-        return pipeline
-    
-
-
-    def give_sprite_scaling_factors(self, sprite_scaling_factors):
-        """
-        Set the sprite scaling factors for the pipeline.
-        This is used to apply per-sprite temperature scaling during inference.
-        """
-        self.sprite_scaling_factors = sprite_scaling_factors
-
-    def __call__(
-        self,
-        batch_size: int = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        num_inference_steps: int = common_settings.NUM_INFERENCE_STEPS,
-        output_type: Optional[str] = "tensor",
-        return_dict: bool = True,
-        height: int = common_settings.MARIO_HEIGHT, width: int = common_settings.MARIO_WIDTH, 
-        latents: Optional[torch.FloatTensor] = None,
-        show_progress_bar=True,
-    ) -> Union[ImagePipelineOutput, Tuple]:
-
-        self.unet.eval()
-        with torch.no_grad():
-
-            if latents is not None:
-                image = latents.to(self.device)
-            else:
-                image_shape = (
-                    batch_size,
-                    self.unet.config.in_channels,
-                    height,
-                    width
-                )
-
-                image = torch.randn(image_shape, generator=generator, device=self.device)
-
-            self.scheduler.set_timesteps(num_inference_steps)
-
-            iterator = self.progress_bar(self.scheduler.timesteps) if show_progress_bar else self.scheduler.timesteps
-            for t in iterator:
-                #print(image.shape)
-                model_output = self.unet(image, t).sample
-                image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
-
-            # Apply per-sprite temperature scaling if enabled
-            if hasattr(self,"sprite_scaling_factors") and self.sprite_scaling_factors is not None:
-                image = image / self.sprite_scaling_factors.view(1, -1, 1, 1)
-
-            
-            if self.block_embeddings is not None:
-                image = get_scene_from_embeddings(image, self.block_embeddings)
-            else:
-                image = F.softmax(image, dim=1)
-                image = image.detach().cpu() 
-
-            if not return_dict:
-                return (image,)
-
-            return ImagePipelineOutput(images=image)
-
-    def print_unet_architecture(self):
-        """Prints the architecture of the UNet model."""
-        print(self.unet)

models/pipeline_loader.py

@@ -1,41 +0,0 @@
-from models.text_diffusion_pipeline import TextConditionalDDPMPipeline
-from models.latent_diffusion_pipeline import UnconditionalDDPMPipeline
-import os
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-
-
-def get_pipeline(model_path):
-    # If model_path is a local directory, use the original logic
-    if os.path.isdir(model_path):
-        #Diffusion models
-        if os.path.exists(os.path.join(model_path, "unet")):
-            if os.path.exists(os.path.join(model_path, "text_encoder")):
-                #If it has a text encoder and a unet, it's text conditional diffusion
-                pipe = TextConditionalDDPMPipeline.from_pretrained(model_path)
-            else:
-                #If it has no text encoder, use the unconditional diffusion model
-                pipe = UnconditionalDDPMPipeline.from_pretrained(model_path)
-    else:
-        # Assume it's a Hugging Face Hub model ID
-        # Try to load config to determine if it's text-conditional
-        try:
-            config, _ = DiffusionPipeline.load_config(model_path)
-            components = config.get("components", {})
-        except Exception:
-            components = {}
-        if "text_encoder" in components or "text_encoder" in str(components):
-            # Use the local pipeline file for custom_pipeline
-            pipe = DiffusionPipeline.from_pretrained(
-                model_path,
-                custom_pipeline="models.text_diffusion_pipeline.TextConditionalDDPMPipeline",
-                trust_remote_code=True,
-            )
-        else:
-            # Fallback: try unconditional 
-            pipe = DiffusionPipeline.from_pretrained(
-                model_path,
-                custom_pipeline="models.latent_diffusion_pipeline.UnconditionalDDPMPipeline",
-                trust_remote_code=True,
-            )
-
-    return pipe

models/sentence_transformers_helper.py

@@ -1,114 +0,0 @@
-from transformers import AutoTokenizer, AutoModel
-import torch
-import torch.nn.functional as F
-
-#Mean Pooling - Take average of all tokens
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output.last_hidden_state
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-
-#Encode text
-def encode(texts, tokenizer, model, device='cpu'):
-    # Tokenize sentences
-    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
-    encoded_input.to(device)
-
-    # Compute token embeddings
-    with torch.no_grad():
-        model_output = model(**encoded_input, return_dict=True)
-
-    # Perform pooling
-    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
-
-    # Normalize embeddings
-    embeddings = F.normalize(embeddings, p=2, dim=1)
-    
-    embeddings = embeddings.to(device)
-
-    return embeddings
-
-# Get embeddings for a batch of captions and optional negative captions
-def get_embeddings(batch_size, tokenizer, model, captions=None, neg_captions=None, device='cpu'):
-    embeddings = encode([""]*batch_size, tokenizer, model, device)
-
-    if captions is not None:
-        caption_embeddings = encode(captions, tokenizer, model, device)
-        embeddings = torch.cat((embeddings, caption_embeddings), dim=0)
-
-    if neg_captions is not None:
-        neg_embeddings = encode(neg_captions, tokenizer, model, device)
-        embeddings = torch.cat((neg_embeddings, embeddings), dim=0)
-    
-    
-    embeddings = embeddings.unsqueeze(1)
-    
-    return embeddings
-
-
-
-
-def get_embeddings_split(batch_size, tokenizer, model, captions=None, neg_captions=None, device='cpu', max_length=20):
-
-    padding_length = max(max([s.count(".") for s in captions]) if captions else 1, 
-                     max([s.count(".") for s in neg_captions]) if neg_captions else 1)
-    if (padding_length>max_length):
-        raise ValueError(f"Token sequence length {padding_length} exceeds specified length {max_length}.")
-
-
-    empty_split = split_sentences([""] * batch_size, padding_length)
-    embeddings = get_embeddings_from_split(empty_split, tokenizer, model, device)
-
-    if(captions is not None):
-        captions_split = split_sentences(captions, padding_length)
-        caption_embeddings = get_embeddings_from_split(captions_split, tokenizer, model, device)
-        embeddings = torch.cat((embeddings, caption_embeddings), dim=0)
-    
-    if(neg_captions is not None):
-        neg_split = split_sentences(neg_captions, padding_length)
-        neg_embeddings = get_embeddings_from_split(neg_split, tokenizer, model, device)
-        embeddings = torch.cat((neg_embeddings, embeddings), dim=0)
-    
-
-    #We don't need to unsqueeze this, we have an array of (batch_size, padding_length, encoding_size) already
-
-    return embeddings.to(device)
-
-
-#This method takes a caption batch in list form, and outputs a 2d list where every caption has been split by period
-def split_sentences(caption_array, padding_length=20):  
-    split_caption_array = []
-
-    #Padding happens here
-    for caption in caption_array:
-        split_caption = [s.strip() for s in caption.split(".") if s.strip()]
-        #This is the token padding, we just use an empty string
-        split_caption += [""] * (padding_length - len(split_caption))
-        split_caption_array.append(split_caption)
-        
-    return split_caption_array
-
-
-#Expects all split vectors to be the same length
-def get_embeddings_from_split(caption_batch, tokenizer, model, device='cpu'):
-    all_caption_encodings = []
-    for caption_sequence in caption_batch:
-        #Encode the sequence of split captions as if it was a batch, should now be a [maxlength, embeddingsize] tensor
-        caption_sequence = encode(caption_sequence, tokenizer, model, device)
-        
-        #We don't reshape this to avoid having to unsqueeze it later
-        all_caption_encodings.append(caption_sequence)
-    
-    all_caption_encodings = torch.stack(all_caption_encodings, dim=0)
-    return all_caption_encodings
-        
-        
-
-if __name__ == "__main__":
-    cap = split_sentences(["Hello. My name is George. How. Are you doing. Today?", "I am doing. Just fine. Thanks."])
-    model_url = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"
-    device = 'cuda'
-    tokenizer = AutoTokenizer.from_pretrained(model_url)
-    model = AutoModel.from_pretrained(model_url, trust_remote_code=True).to(device)
-    get_embeddings_from_split(cap, tokenizer, model, device)

models/text_diffusion_pipeline.py

@@ -1,442 +0,0 @@
-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

@@ -1,206 +0,0 @@
-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}")