main.py

from dataclasses import dataclass
from enum import Enum, auto
from pathlib import Path
from typing import List, Optional, Union, Set, Dict, Any, Tuple
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers import FluxPipeline, MochiPipeline
from loguru import logger
from timm import create_model
from transformers import AutoTokenizer, AutoModel

class ModalityType(Enum):
    TEXT = auto()
    IMAGE = auto()
    VIDEO = auto()
    AUDIO = auto()

@dataclass
class ModalityConfig:
    """Configuration for individual modality."""
    enabled: bool = True
    model_path: str = ""
    embedding_dim: int = 1024
    max_sequence_length: int = 512

@dataclass
class ModelConfig:
    """Configuration for the Dynamic Multi-Modal Model."""
    modalities: Dict[ModalityType, ModalityConfig] = None
    fusion_dim: int = 1024
    num_fusion_layers: int = 4
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    torch_dtype: torch.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    
    def __post_init__(self):
        if self.modalities is None:
            self.modalities = {
                ModalityType.TEXT: ModalityConfig(model_path="dfurman/CalmeRys-78B-Orpo-v0.1"),
                ModalityType.IMAGE: ModalityConfig(model_path="black-forest-labs/FLUX.1-dev"),
                ModalityType.VIDEO: ModalityConfig(model_path="genmo/mochi-1-preview"),
            }

class ModalityFusion(nn.Module):
    """Cross-attention based fusion of different modality embeddings."""
    
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        # Multi-head cross-attention layers
        self.fusion_layers = nn.ModuleList([
            nn.MultiheadAttention(
                config.fusion_dim,
                num_heads=8,
                batch_first=True
            ) for _ in range(config.num_fusion_layers)
        ])
        
        # Modality-specific projections
        self.modality_projections = nn.ModuleDict({
            modality.name.lower(): nn.Linear(
                config.modalities[modality].embedding_dim,
                config.fusion_dim
            ) for modality in config.modalities
        })
        
        # Output modality classifier
        self.modality_classifier = nn.Linear(
            config.fusion_dim,
            len(ModalityType)
        )
        
    def forward(
        self,
        embeddings: Dict[ModalityType, torch.Tensor]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Fuse multiple modality embeddings and predict output modalities.
        
        Args:
            embeddings: Dictionary of embeddings per modality
            
        Returns:
            Tuple of (fused embeddings, output modality probabilities)
        """
        # Project each modality to common dimension
        projected_embeddings = []
        for modality, embedding in embeddings.items():
            proj = self.modality_projections[modality.name.lower()]
            projected_embeddings.append(proj(embedding))
        
        # Concatenate all embeddings
        fused = torch.cat(projected_embeddings, dim=1)
        
        # Apply fusion layers
        for layer in self.fusion_layers:
            fused_attn, _ = layer(fused, fused, fused)
            fused = fused + fused_attn
        
        # Predict output modalities
        modality_logits = self.modality_classifier(fused.mean(dim=1))
        modality_probs = torch.sigmoid(modality_logits)
        
        return fused, modality_probs

class DynamicMultiModal:
    """Dynamic multi-modal model with automatic modality selection."""
    
    def __init__(self, config: ModelConfig):
        logger.info("Initializing Dynamic Multi-Modal Model")
        self.config = config
        
        # Initialize modality-specific models
        self._init_models()
        
        # Initialize fusion module
        self.fusion_module = ModalityFusion(config).to(
            config.device,
            dtype=config.torch_dtype
        )
        
        logger.info("Model initialization complete")
        
    def _init_models(self):
        """Initialize all modality-specific models."""
        self.models = {}
        self.processors = {}
        
        for modality, mod_config in self.config.modalities.items():
            if not mod_config.enabled:
                continue
                
            logger.info(f"Initializing {modality.name} model")
            
            if modality == ModalityType.TEXT:
                self.processors[modality] = AutoTokenizer.from_pretrained(
                    mod_config.model_path
                )
                self.models[modality] = AutoModel.from_pretrained(
                    mod_config.model_path,
                    torch_dtype=self.config.torch_dtype,
                    device_map="auto"
                )
                
            elif modality == ModalityType.IMAGE:
                self.models[modality] = FluxPipeline.from_pretrained(
                    mod_config.model_path,
                    torch_dtype=self.config.torch_dtype
                )
                self.models[modality].enable_model_cpu_offload()
                
            elif modality == ModalityType.VIDEO:
                self.models[modality] = MochiPipeline.from_pretrained(
                    mod_config.model_path,
                    variant="bf16",
                    torch_dtype=self.config.torch_dtype
                )
                self.models[modality].enable_model_cpu_offload()
                self.models[modality].enable_vae_tiling()

    def process_inputs(
        self,
        inputs: Dict[ModalityType, Union[str, Path, torch.Tensor]]
    ) -> Dict[ModalityType, torch.Tensor]:
        """Process multiple input modalities into embeddings."""
        embeddings = {}
        
        for modality, input_data in inputs.items():
            if modality not in self.config.modalities or \
               not self.config.modalities[modality].enabled:
                continue
                
            embeddings[modality] = self._process_modality(modality, input_data)
            
        return embeddings

    def _process_modality(
        self,
        modality: ModalityType,
        input_data: Union[str, Path, torch.Tensor]
    ) -> torch.Tensor:
        """Process single modality input into embeddings."""
        if modality == ModalityType.TEXT:
            inputs = self.processors[modality](
                input_data,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.config.modalities[modality].max_sequence_length
            ).to(self.config.device)
            
            with torch.no_grad():
                outputs = self.models[modality](**inputs)
            return outputs.last_hidden_state
            
        elif modality == ModalityType.IMAGE:
            # Process image through vision encoder
            if isinstance(input_data, str):
                input_data = Path(input_data)
            # Return image features
            return self.models[modality].encode_image(input_data)
            
        raise ValueError(f"Unsupported input modality: {modality}")

    def generate(
        self,
        inputs: Dict[ModalityType, Union[str, Path, torch.Tensor]],
        prompt: Optional[str] = None,
        force_modalities: Optional[Set[ModalityType]] = None,
        **kwargs
    ) -> Dict[ModalityType, torch.Tensor]:
        """
        Generate outputs in automatically determined modalities.
        
        Args:
            inputs: Dictionary of input data per modality
            prompt: Optional text prompt for generation
            force_modalities: Optional set of modalities to force generate
            **kwargs: Additional generation parameters
            
        Returns:
            Dictionary of generated outputs per modality
        """
        logger.info("Processing inputs")
        embeddings = self.process_inputs(inputs)
        
        # Fuse modalities and predict output types
        logger.info("Fusing modalities")
        fused_embeddings, modality_probs = self.fusion_module(embeddings)
        
        # Determine output modalities
        if force_modalities:
            output_modalities = force_modalities
        else:
            # Select modalities above threshold
            output_modalities = {
                modality for i, modality in enumerate(ModalityType)
                if modality_probs[0, i] > 0.5
            }
        
        logger.info(f"Generating outputs for modalities: {output_modalities}")
        
        outputs = {}
        for modality in output_modalities:
            if modality == ModalityType.IMAGE:
                outputs[modality] = self._generate_image(
                    fused_embeddings,
                    prompt,
                    **kwargs
                )
            elif modality == ModalityType.VIDEO:
                outputs[modality] = self._generate_video(
                    fused_embeddings,
                    prompt,
                    **kwargs
                )
                
        return outputs

    def _generate_image(
        self,
        embeddings: torch.Tensor,
        prompt: str,
        height: int = 1024,
        width: int = 1024,
        **kwargs
    ) -> torch.Tensor:
        """Generate image output."""
        image = self.models[ModalityType.IMAGE](
            prompt,
            height=height,
            width=width,
            guidance_scale=kwargs.get('guidance_scale', 3.5),
            num_inference_steps=kwargs.get('num_inference_steps', 50),
            max_sequence_length=kwargs.get('max_sequence_length', 512),
            generator=torch.Generator(self.config.device).manual_seed(
                kwargs.get('seed', 0)
            )
        ).images[0]
        
        return torch.tensor(image)

    def _generate_video(
        self,
        embeddings: torch.Tensor,
        prompt: str,
        num_frames: int = 84,
        **kwargs
    ) -> torch.Tensor:
        """Generate video output."""
        frames = self.models[ModalityType.VIDEO](
            prompt,
            num_frames=num_frames,
            **kwargs
        ).frames[0]
        
        return torch.tensor(frames)

def setup_logging(log_file: Optional[str] = None):
    """Configure logging with loguru."""
    logger.remove()
    
    logger.add(
        lambda msg: print(msg),
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <cyan>{function}</cyan>: <white>{message}</white>",
        level="INFO"
    )
    
    if log_file:
        logger.add(
            log_file,
            rotation="500 MB",
            retention="10 days",
            compression="zip",
            level="DEBUG"
        )

def main():
    """Example usage of the Dynamic Multi-Modal Model."""
    setup_logging("multimodal.log")
    
    config = ModelConfig()
    model = DynamicMultiModal(config)
    
    # Example: Multiple inputs, automatic output selection
    inputs = {
        ModalityType.TEXT: "A beautiful sunset over mountains",
        ModalityType.IMAGE: "path/to/reference_image.jpg"
    }
    
    # Generate outputs - model will automatically determine modalities
    outputs = model.generate(inputs, prompt="Mountain sunset scene")
    
    # Force specific output modalities
    outputs = model.generate(
        inputs,
        prompt="Mountain sunset scene",
        force_modalities={ModalityType.IMAGE, ModalityType.VIDEO}
    )
    
    logger.info("Generation complete")

if __name__ == "__main__":
    main()