""" DNA-Diffusion Model Wrapper Singleton class to handle model loading and sequence generation """ import os import sys import torch import numpy as np import logging from typing import Optional, Dict, List import time logger = logging.getLogger(__name__) class DNADiffusionModel: """Singleton wrapper for DNA-Diffusion model""" _instance = None _initialized = False # Cell type mapping from simple names to dataset identifiers CELL_TYPE_MAPPING = { 'K562': 'K562_ENCLB843GMH', 'GM12878': 'GM12878_ENCLB441ZZZ', 'HepG2': 'HepG2_ENCLB029COU', 'hESCT0': 'hESCT0_ENCLB449ZZZ' } def __new__(cls): if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def __init__(self): """Initialize the model (only runs once due to singleton pattern)""" if not self._initialized: self._initialize() self._initialized = True def _initialize(self): """Load model and setup components""" try: logger.info("Initializing DNA-Diffusion model...") # Add DNA-Diffusion to path dna_diffusion_path = os.path.join(os.path.dirname(__file__), 'DNA-Diffusion') if os.path.exists(dna_diffusion_path): sys.path.insert(0, os.path.join(dna_diffusion_path, 'src')) # Import DNA-Diffusion components from dnadiffusion.models.pretrained_unet import PretrainedUNet from dnadiffusion.models.diffusion import Diffusion from dnadiffusion.data.dataloader import get_dataset_for_sampling # Load pretrained model from HuggingFace logger.info("Loading pretrained model from HuggingFace...") self.model = PretrainedUNet.from_pretrained("ssenan/DNA-Diffusion") # Move to GPU if available self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"Using device: {self.device}") self.model = self.model.to(self.device) self.model.eval() # Initialize diffusion sampler with the model self.diffusion = Diffusion( model=self.model, timesteps=50, beta_start=0.0001, beta_end=0.2 ) # Ensure output_attention is set to False initially if hasattr(self.model, 'output_attention'): self.model.output_attention = False if hasattr(self.model.model, 'output_attention'): self.model.model.output_attention = False # Setup dataset for sampling data_path = os.path.join(dna_diffusion_path, "data/K562_hESCT0_HepG2_GM12878_12k_sequences_per_group.txt") saved_data_path = os.path.join(dna_diffusion_path, "data/encode_data.pkl") # Get dataset info train_data, val_data, cell_num_list, numeric_to_tag_dict = get_dataset_for_sampling( data_path=data_path, saved_data_path=saved_data_path, load_saved_data=True, debug=False, cell_types=None # Load all cell types ) # Store dataset info self.train_data = train_data self.val_data = val_data self.cell_num_list = cell_num_list self.numeric_to_tag_dict = numeric_to_tag_dict # Get available cell types self.available_cell_types = [numeric_to_tag_dict[num] for num in cell_num_list] logger.info(f"Available cell types: {self.available_cell_types}") # Warm up the model with a test generation logger.info("Warming up model...") self._warmup() logger.info("Model initialization complete!") except Exception as e: logger.error(f"Failed to initialize model: {str(e)}") self.model = None self.diffusion = None self.dataset = None raise def _warmup(self): """Warm up the model with a test generation""" try: # Generate one sequence for the first available cell type if self.available_cell_types: cell_type = list(self.CELL_TYPE_MAPPING.keys())[0] self.generate(cell_type, guidance_scale=1.0) except Exception as e: logger.warning(f"Warmup generation failed: {str(e)}") def is_ready(self) -> bool: """Check if model is loaded and ready""" return self.model is not None and self.diffusion is not None and self.train_data is not None def generate(self, cell_type: str, guidance_scale: float = 1.0) -> Dict[str, any]: """ Generate a DNA sequence for the specified cell type Args: cell_type: Simple cell type name (K562, GM12878, HepG2, hESCT0) guidance_scale: Guidance scale for generation (1.0-10.0) Returns: Dict with 'sequence' (200bp string) and 'metadata' """ if not self.is_ready(): raise RuntimeError("Model is not initialized") # Validate inputs if cell_type not in self.CELL_TYPE_MAPPING: raise ValueError(f"Invalid cell type: {cell_type}. Must be one of {list(self.CELL_TYPE_MAPPING.keys())}") if not 1.0 <= guidance_scale <= 10.0: raise ValueError(f"Guidance scale must be between 1.0 and 10.0, got {guidance_scale}") # Map to full cell type identifier full_cell_type = self.CELL_TYPE_MAPPING[cell_type] # Find the numeric index for this cell type tag_to_numeric = {tag: num for num, tag in self.numeric_to_tag_dict.items()} # Find matching cell type (case-insensitive partial match) cell_type_numeric = None for tag, num in tag_to_numeric.items(): if full_cell_type.lower() in tag.lower() or tag.lower() in full_cell_type.lower(): cell_type_numeric = num logger.info(f"Matched '{full_cell_type}' to '{tag}'") break if cell_type_numeric is None: raise ValueError(f"Cell type {full_cell_type} not found in dataset. Available: {list(self.numeric_to_tag_dict.values())}") try: logger.info(f"Generating sequence for {cell_type} (guidance={guidance_scale})...") start_time = time.time() # For now, use simple generation without classifier-free guidance # TODO: Fix classifier-free guidance implementation sequence = self._generate_simple(cell_type_numeric, guidance_scale) generation_time = time.time() - start_time logger.info(f"Generated sequence in {generation_time:.2f}s") return { 'sequence': sequence, 'metadata': { 'cell_type': cell_type, 'full_cell_type': full_cell_type, 'guidance_scale': guidance_scale, 'generation_time': generation_time, 'sequence_length': len(sequence) } } except Exception as e: logger.error(f"Generation failed: {str(e)}") raise def _generate_simple(self, cell_type_idx: int, guidance_scale: float) -> str: """Simple generation using the diffusion model's sample method""" with torch.no_grad(): # For guidance_scale = 1.0, use simple generation without classifier-free guidance if guidance_scale == 1.0: # Create initial noise img = torch.randn((1, 1, 4, 200), device=self.device) # Simple denoising loop without guidance for i in reversed(range(self.diffusion.timesteps)): t = torch.full((1,), i, device=self.device, dtype=torch.long) # Get model prediction with classes classes = torch.tensor([cell_type_idx], device=self.device, dtype=torch.long) noise_pred = self.model(img, time=t, classes=classes) # Denoising step betas_t = self.diffusion.betas[i] sqrt_one_minus_alphas_cumprod_t = self.diffusion.sqrt_one_minus_alphas_cumprod[i] sqrt_recip_alphas_t = self.diffusion.sqrt_recip_alphas[i] # Predict x0 model_mean = sqrt_recip_alphas_t * (img - betas_t * noise_pred / sqrt_one_minus_alphas_cumprod_t) if i == 0: img = model_mean else: posterior_variance_t = self.diffusion.posterior_variance[i] noise = torch.randn_like(img) img = model_mean + torch.sqrt(posterior_variance_t) * noise final_image = img[0] # Remove batch dimension else: # Use the diffusion model's built-in sample method with guidance # This requires proper context mask handling which is complex # For now, fall back to simple generation logger.warning(f"Guidance scale {guidance_scale} not fully implemented, using simple generation") return self._generate_simple(cell_type_idx, 1.0) # Convert to sequence final_array = final_image.cpu().numpy() sequence = self._array_to_sequence(final_array) return sequence def _array_to_sequence(self, array: np.ndarray) -> str: """Convert model output array to DNA sequence string""" # Get nucleotide mapping nucleotides = ['A', 'C', 'G', 'T'] # array shape is (1, 4, 200) - channels, nucleotides, sequence_length # Reshape to (4, 200) and get argmax along nucleotide dimension array = array.squeeze(0) # Remove channel dimension -> (4, 200) indices = np.argmax(array, axis=0) # Get max nucleotide for each position # Convert indices to nucleotides sequence = ''.join(nucleotides[int(idx)] for idx in indices) return sequence def get_model_info(self) -> Dict[str, any]: """Get information about the loaded model""" if not self.is_ready(): return {'status': 'not_initialized'} return { 'status': 'ready', 'device': str(self.device), 'cell_types': list(self.CELL_TYPE_MAPPING.keys()), 'full_cell_types': self.available_cell_types, 'model_name': 'ssenan/DNA-Diffusion', 'sequence_length': 200, 'guidance_scale_range': [1.0, 10.0] } # Convenience functions for direct usage _model_instance = None def get_model() -> DNADiffusionModel: """Get or create the singleton model instance""" global _model_instance if _model_instance is None: _model_instance = DNADiffusionModel() return _model_instance def generate_sequence(cell_type: str, guidance_scale: float = 1.0) -> str: """Generate a DNA sequence (convenience function)""" model = get_model() result = model.generate(cell_type, guidance_scale) return result['sequence']