Add model

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,98 @@
 ---
 license: mit
+library_name: diffusers
+pipeline_tag: text-to-image
+language:
+- en
+tags:
+- diffusers
+- matfuse
+- pbr
+- material-generation
+- svbrdf
+- text-to-image
 ---
+
+
+# MatFuse — Controllable Material Generation with Diffusion Models
+
+MatFuse generates tileable PBR material maps (diffuse, normal, roughness,
+specular) from text, reference images, sketches, and/or color palettes.
+
+> **Paper:** [MatFuse: Controllable Material Generation with Diffusion Models](https://arxiv.org/abs/2308.11408) — CVPR 2024
+> **Project page:** <https://gvecchio.com/matfuse/>
+
+## Quick Start
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "gvecchio/MatFuse",
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+)
+pipe = pipe.to("cuda")
+
+result = pipe(
+    text="red brick wall",
+    num_inference_steps=50,
+    guidance_scale=4.0,
+    generator=torch.Generator("cuda").manual_seed(42),
+)
+
+result["diffuse"][0].save("diffuse.png")
+result["normal"][0].save("normal.png")
+result["roughness"][0].save("roughness.png")
+result["specular"][0].save("specular.png")
+```
+
+## Conditioning Inputs
+
+All conditions are **optional** and freely composable:
+
+| Input | Type | Description |
+|-------|------|-------------|
+| `text` | `str` | Text description of the material |
+| `image` | `PIL.Image` | Reference image for style/appearance |
+| `sketch` | `PIL.Image` (grayscale) | Binary edge map for structure |
+| `palette` | `list[tuple]` | Up to 5 RGB colour tuples (0–255) |
+
+```python
+from PIL import Image
+
+result = pipe(
+    image=Image.open("reference.png"),
+    text="rough stone texture",
+    palette=[(120, 80, 60), (90, 60, 40), (150, 110, 80), (70, 50, 30), (180, 140, 100)],
+    num_inference_steps=50,
+    guidance_scale=4.0,
+)
+```
+
+## Architecture
+
+| Component | Class | Key parameters |
+|-----------|-------|----------------|
+| **UNet** | `UNet2DConditionModel` | in=16, out=12, blocks=[256,512,1024], cross_attn=512 |
+| **VAE** | `MatFuseVQModel` (custom) | 4 encoders + 4 VQ codebooks (4096×3), shared decoder, f=8 |
+| **Scheduler** | `DDIMScheduler` | β 0.0015–0.0195, scaled_linear, ε-prediction |
+| **Conditioning** | `MultiConditionEncoder` (custom) | CLIP ViT-B/16 · sentence-transformers · palette MLP · sketch CNN |
+
+## 📜 Citation
+
+```bibtex
+@inproceedings{vecchio2024matfuse,
+  author    = {Vecchio, Giuseppe and Sortino, Renato and Palazzo, Simone and Spampinato, Concetto},
+  title     = {MatFuse: Controllable Material Generation with Diffusion Models},
+  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  month     = {June},
+  year      = {2024},
+  pages     = {4429-4438}
+}
+```
+
+## License
+
+This project is licensed under the MIT License.

checkpoints/matfuse-full.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0f9ce7058f0a122dcbca8d4870d98df656c8114f20914e7d126406222b5df98
+size 4641175367

checkpoints/matfuse-pruned.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27597775fd7f25a98523e7171bf31305bc3f274f7300fbd406654952e608a036
+size 3060674264

condition_encoder/condition_encoders.py

@@ -0,0 +1,426 @@
+"""
+MatFuse Condition Encoders for diffusers.
+
+These encoders handle the multi-modal conditioning:
+- Image embedding (CLIP image encoder)
+- Text embedding (CLIP text encoder)
+- Sketch encoder (CNN)
+- Palette encoder (MLP)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Dict, Union, List
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+class SketchEncoder(ModelMixin, ConfigMixin):
+    """
+    CNN encoder for binary sketch/edge maps.
+
+    Takes a single-channel binary image and encodes it to a spatial feature map
+    that will be concatenated with the latent for hybrid conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 1,
+        out_channels: int = 4,
+    ):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, 32, 7, 1, 1),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.Conv2d(32, 64, 3, 2, 1),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.Conv2d(64, 128, 3, 2, 1),
+            nn.BatchNorm2d(128),
+            nn.GELU(),
+            nn.Conv2d(128, 256, 3, 2, 1),
+            nn.BatchNorm2d(256),
+            nn.GELU(),
+            nn.Conv2d(256, out_channels, 1, 1, 0),
+            nn.BatchNorm2d(out_channels),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode sketch input.
+
+        Args:
+            x: Input tensor of shape (B, 1, H, W) with values in [0, 1].
+
+        Returns:
+            Encoded features of shape (B, out_channels, H/8, W/8).
+        """
+        return self.net(x)
+
+
+class PaletteEncoder(ModelMixin, ConfigMixin):
+    """
+    MLP encoder for color palettes.
+
+    Takes a color palette (N colors, RGB) and encodes it to a single embedding
+    for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        hidden_channels: int = 64,
+        out_channels: int = 512,
+        n_colors: int = 5,
+    ):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(in_channels, hidden_channels),
+            nn.GELU(),
+            nn.Flatten(),
+            nn.Linear(hidden_channels * n_colors, out_channels),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode color palette.
+
+        Args:
+            x: Input tensor of shape (B, n_colors, 3) with RGB values in [0, 1].
+
+        Returns:
+            Encoded embedding of shape (B, out_channels).
+        """
+        return self.net(x)
+
+
+class CLIPImageEncoder(ModelMixin, ConfigMixin):
+    """
+    Wrapper for CLIP image encoder using the OpenAI CLIP library.
+
+    Generates image embeddings for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_name: str = "ViT-B/16",
+        normalize: bool = True,
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.normalize = normalize
+        self.model = None  # Lazy loading
+
+        # Register normalization buffers
+        self.register_buffer(
+            "mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False
+        )
+        self.register_buffer(
+            "std", torch.tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False
+        )
+
+    def _load_model(self):
+        """Lazy load the CLIP model."""
+        if self.model is None:
+            import clip
+
+            self.model, _ = clip.load(self.model_name, device="cpu", jit=False)
+            self.model = self.model.visual
+
+    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
+        """Preprocess images for CLIP."""
+        # Resize to 224x224
+        x = F.interpolate(
+            x, size=(224, 224), mode="bicubic", align_corners=True, antialias=True
+        )
+        # Normalize from [-1, 1] to [0, 1]
+        x = (x + 1.0) / 2.0
+        # Normalize according to CLIP - move mean/std to device if needed
+        mean = self.mean.to(x.device).view(1, 3, 1, 1)
+        std = self.std.to(x.device).view(1, 3, 1, 1)
+        x = (x - mean) / std
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode image using CLIP.
+
+        Args:
+            x: Input tensor of shape (B, 3, H, W) with values in [-1, 1].
+
+        Returns:
+            Image embedding of shape (B, 1, 512).
+        """
+        self._load_model()
+
+        # Move model to same device as input
+        device = x.device
+        self.model = self.model.to(device)
+
+        x = self.preprocess(x)
+        z = self.model(x).float().unsqueeze(1)  # (B, 1, 512)
+
+        if self.normalize:
+            z = z / torch.linalg.norm(z, dim=2, keepdim=True)
+
+        return z
+
+
+class CLIPTextEncoder(ModelMixin, ConfigMixin):
+    """
+    Wrapper for CLIP sentence encoder using sentence-transformers.
+
+    Generates text embeddings for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/clip-ViT-B-16",
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.model = None  # Lazy loading
+
+    def _load_model(self):
+        """Lazy load the sentence transformer model."""
+        if self.model is None:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model_name)
+            self.model.eval()
+
+    def forward(self, text: Union[str, List[str]]) -> torch.Tensor:
+        """
+        Encode text using CLIP sentence transformer.
+
+        Args:
+            text: Input text or list of texts.
+
+        Returns:
+            Text embedding of shape (B, 512).
+        """
+        self._load_model()
+
+        if isinstance(text, str):
+            text = [text]
+
+        embeddings = self.model.encode(text, convert_to_tensor=True)
+        return embeddings
+
+
+class MultiConditionEncoder(ModelMixin, ConfigMixin):
+    """
+    Multi-condition encoder that combines all conditioning modalities.
+
+    This encoder takes multiple condition inputs and produces:
+    - c_crossattn: Features for cross-attention (image, text, palette embeddings)
+    - c_concat: Features for concatenation (sketch encoding)
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sketch_in_channels: int = 1,
+        sketch_out_channels: int = 4,
+        palette_in_channels: int = 3,
+        palette_hidden_channels: int = 64,
+        palette_out_channels: int = 512,
+        n_colors: int = 5,
+        clip_image_model: str = "ViT-B/16",
+        clip_text_model: str = "sentence-transformers/clip-ViT-B-16",
+    ):
+        super().__init__()
+
+        self.sketch_encoder = SketchEncoder(
+            in_channels=sketch_in_channels,
+            out_channels=sketch_out_channels,
+        )
+
+        self.palette_encoder = PaletteEncoder(
+            in_channels=palette_in_channels,
+            hidden_channels=palette_hidden_channels,
+            out_channels=palette_out_channels,
+            n_colors=n_colors,
+        )
+
+        # CLIP encoders are lazy-loaded
+        self.clip_image_encoder = None
+        self.clip_text_encoder = None
+        self._clip_image_model = clip_image_model
+        self._clip_text_model = clip_text_model
+
+    def _load_clip_encoders(self):
+        """Lazy load CLIP encoders."""
+        if self.clip_image_encoder is None:
+            self.clip_image_encoder = CLIPImageEncoder(
+                model_name=self._clip_image_model
+            )
+        if self.clip_text_encoder is None:
+            self.clip_text_encoder = CLIPTextEncoder(model_name=self._clip_text_model)
+
+    def encode_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Encode image using CLIP."""
+        self._load_clip_encoders()
+        return self.clip_image_encoder(image)
+
+    def encode_text(self, text: Union[str, List[str]]) -> torch.Tensor:
+        """Encode text using CLIP."""
+        self._load_clip_encoders()
+        return self.clip_text_encoder(text)
+
+    def encode_sketch(self, sketch: torch.Tensor) -> torch.Tensor:
+        """Encode sketch/edge map."""
+        return self.sketch_encoder(sketch)
+
+    def encode_palette(self, palette: torch.Tensor) -> torch.Tensor:
+        """Encode color palette."""
+        return self.palette_encoder(palette)
+
+    def get_unconditional_conditioning(
+        self,
+        batch_size: int = 1,
+        image_size: int = 256,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Get unconditional conditioning for classifier-free guidance.
+
+        IMPORTANT: The original model was trained to drop conditions by replacing them
+        with encoded placeholders (zero/gray image through CLIP, empty string through
+        sentence-transformers, zero palette through PaletteEncoder, zero sketch through
+        SketchEncoder) — NOT with zero tensors. This method produces the correct
+        unconditional embeddings.
+
+        Args:
+            batch_size: Batch size.
+            image_size: Image resolution (for sketch spatial dims).
+            device: Device to place tensors on.
+
+        Returns:
+            Dictionary with c_crossattn and c_concat for unconditional guidance.
+        """
+        return self.forward(
+            image_embed=None,
+            text=None,
+            sketch=None,
+            palette=None,
+            batch_size=batch_size,
+            image_size=image_size,
+            device=device,
+        )
+
+    def forward(
+        self,
+        image_embed: Optional[torch.Tensor] = None,
+        text: Optional[Union[str, List[str]]] = None,
+        sketch: Optional[torch.Tensor] = None,
+        palette: Optional[torch.Tensor] = None,
+        batch_size: int = 1,
+        image_size: int = 256,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Encode all conditions.
+
+        When a condition is not provided, the model encodes a placeholder input
+        through the actual encoder (matching training behavior) rather than using
+        zero tensors. This is critical because the model was trained with:
+        - Image drop → CLIP encoding of a gray/zero image (0.0 in [-1,1])
+        - Text drop → sentence-transformer encoding of ""
+        - Palette drop → PaletteEncoder(zeros)
+        - Sketch drop → SketchEncoder(zeros)
+
+        Args:
+            image_embed: Reference image of shape (B, 3, H, W) in [-1, 1].
+            text: Text description(s).
+            sketch: Binary sketch of shape (B, 1, H, W) in [0, 1].
+            palette: Color palette of shape (B, n_colors, 3) in [0, 1].
+            batch_size: Batch size (used when no inputs are provided).
+            image_size: Image resolution (used to create placeholder sketch).
+            device: Device to place tensors on.
+
+        Returns:
+            Dictionary with:
+            - c_crossattn: Cross-attention context of shape (B, 3, 512) - always 3 tokens.
+            - c_concat: Concatenation features of shape (B, 4, H/8, W/8).
+        """
+        self._load_clip_encoders()
+
+        # Determine batch size and device from any available input
+        if image_embed is not None:
+            batch_size = image_embed.shape[0]
+            device = device or image_embed.device
+            image_size = image_embed.shape[-1]
+        elif sketch is not None:
+            batch_size = sketch.shape[0]
+            device = device or sketch.device
+            image_size = sketch.shape[-1]
+        elif palette is not None:
+            batch_size = palette.shape[0]
+            device = device or palette.device
+
+        device = device or torch.device("cpu")
+        # Infer dtype from model weights for placeholder tensors (e.g. float16)
+        dtype = next(self.palette_encoder.parameters()).dtype
+
+        # --- Image embedding (token 0) ---
+        # When not provided, encode a zero (gray) image through CLIP, matching training ucg_training val=0.0
+        if image_embed is not None:
+            img_emb = self.clip_image_encoder(image_embed)  # (B, 1, 512)
+        else:
+            placeholder_img = torch.zeros(
+                batch_size, 3, image_size, image_size, device=device, dtype=dtype
+            )
+            img_emb = self.clip_image_encoder(placeholder_img)  # (B, 1, 512)
+
+        # --- Text embedding (token 1) ---
+        # When not provided, encode empty string through sentence-transformers, matching training ucg_training val=""
+        if text is not None:
+            text_emb = self.clip_text_encoder(text)  # (B, 512)
+            if device is not None:
+                text_emb = text_emb.to(device)
+            text_emb = text_emb.unsqueeze(1)  # (B, 1, 512)
+        else:
+            text_emb = self.clip_text_encoder([""] * batch_size)  # (B, 512)
+            text_emb = text_emb.to(device).unsqueeze(1)  # (B, 1, 512)
+
+        # --- Palette embedding (token 2) ---
+        # When not provided, encode zero palette through PaletteEncoder, matching training ucg_training val=0.0
+        if palette is not None:
+            palette_emb = self.palette_encoder(palette)  # (B, 512)
+            palette_emb = palette_emb.unsqueeze(1)  # (B, 1, 512)
+        else:
+            n_colors = self.config.get("n_colors", 5)
+            placeholder_palette = torch.zeros(batch_size, n_colors, 3, device=device, dtype=dtype)
+            palette_emb = self.palette_encoder(placeholder_palette)  # (B, 512)
+            palette_emb = palette_emb.unsqueeze(1)  # (B, 1, 512)
+
+        # Combine cross-attention embeddings - always (B, 3, 512)
+        c_crossattn = torch.cat([img_emb, text_emb, palette_emb], dim=1)
+
+        # --- Sketch encoding for concatenation ---
+        # When not provided, encode zero sketch through SketchEncoder, matching training ucg_training val=0.0
+        if sketch is not None:
+            c_concat = self.sketch_encoder(sketch)  # (B, 4, H/8, W/8)
+        else:
+            placeholder_sketch = torch.zeros(
+                batch_size, 1, image_size, image_size, device=device, dtype=dtype
+            )
+            c_concat = self.sketch_encoder(placeholder_sketch)  # (B, 4, H/8, W/8)
+
+        return {
+            "c_crossattn": c_crossattn,
+            "c_concat": c_concat,
+        }

condition_encoder/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "MultiConditionEncoder",
+  "_diffusers_version": "0.35.2",
+  "clip_image_model": "ViT-B/16",
+  "clip_text_model": "sentence-transformers/clip-ViT-B-16",
+  "n_colors": 5,
+  "palette_hidden_channels": 64,
+  "palette_in_channels": 3,
+  "palette_out_channels": 512,
+  "sketch_in_channels": 1,
+  "sketch_out_channels": 4
+}

condition_encoder/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bf932008656551c46a064549570965c19c48539c9e6d5a54c2a40a6414518cb
+size 2230464

condition_encoders.py

@@ -0,0 +1,426 @@
+"""
+MatFuse Condition Encoders for diffusers.
+
+These encoders handle the multi-modal conditioning:
+- Image embedding (CLIP image encoder)
+- Text embedding (CLIP text encoder)
+- Sketch encoder (CNN)
+- Palette encoder (MLP)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Dict, Union, List
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+class SketchEncoder(ModelMixin, ConfigMixin):
+    """
+    CNN encoder for binary sketch/edge maps.
+
+    Takes a single-channel binary image and encodes it to a spatial feature map
+    that will be concatenated with the latent for hybrid conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 1,
+        out_channels: int = 4,
+    ):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, 32, 7, 1, 1),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.Conv2d(32, 64, 3, 2, 1),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.Conv2d(64, 128, 3, 2, 1),
+            nn.BatchNorm2d(128),
+            nn.GELU(),
+            nn.Conv2d(128, 256, 3, 2, 1),
+            nn.BatchNorm2d(256),
+            nn.GELU(),
+            nn.Conv2d(256, out_channels, 1, 1, 0),
+            nn.BatchNorm2d(out_channels),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode sketch input.
+
+        Args:
+            x: Input tensor of shape (B, 1, H, W) with values in [0, 1].
+
+        Returns:
+            Encoded features of shape (B, out_channels, H/8, W/8).
+        """
+        return self.net(x)
+
+
+class PaletteEncoder(ModelMixin, ConfigMixin):
+    """
+    MLP encoder for color palettes.
+
+    Takes a color palette (N colors, RGB) and encodes it to a single embedding
+    for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        hidden_channels: int = 64,
+        out_channels: int = 512,
+        n_colors: int = 5,
+    ):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(in_channels, hidden_channels),
+            nn.GELU(),
+            nn.Flatten(),
+            nn.Linear(hidden_channels * n_colors, out_channels),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode color palette.
+
+        Args:
+            x: Input tensor of shape (B, n_colors, 3) with RGB values in [0, 1].
+
+        Returns:
+            Encoded embedding of shape (B, out_channels).
+        """
+        return self.net(x)
+
+
+class CLIPImageEncoder(ModelMixin, ConfigMixin):
+    """
+    Wrapper for CLIP image encoder using the OpenAI CLIP library.
+
+    Generates image embeddings for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_name: str = "ViT-B/16",
+        normalize: bool = True,
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.normalize = normalize
+        self.model = None  # Lazy loading
+
+        # Register normalization buffers
+        self.register_buffer(
+            "mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False
+        )
+        self.register_buffer(
+            "std", torch.tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False
+        )
+
+    def _load_model(self):
+        """Lazy load the CLIP model."""
+        if self.model is None:
+            import clip
+
+            self.model, _ = clip.load(self.model_name, device="cpu", jit=False)
+            self.model = self.model.visual
+
+    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
+        """Preprocess images for CLIP."""
+        # Resize to 224x224
+        x = F.interpolate(
+            x, size=(224, 224), mode="bicubic", align_corners=True, antialias=True
+        )
+        # Normalize from [-1, 1] to [0, 1]
+        x = (x + 1.0) / 2.0
+        # Normalize according to CLIP - move mean/std to device if needed
+        mean = self.mean.to(x.device).view(1, 3, 1, 1)
+        std = self.std.to(x.device).view(1, 3, 1, 1)
+        x = (x - mean) / std
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode image using CLIP.
+
+        Args:
+            x: Input tensor of shape (B, 3, H, W) with values in [-1, 1].
+
+        Returns:
+            Image embedding of shape (B, 1, 512).
+        """
+        self._load_model()
+
+        # Move model to same device as input
+        device = x.device
+        self.model = self.model.to(device)
+
+        x = self.preprocess(x)
+        z = self.model(x).float().unsqueeze(1)  # (B, 1, 512)
+
+        if self.normalize:
+            z = z / torch.linalg.norm(z, dim=2, keepdim=True)
+
+        return z
+
+
+class CLIPTextEncoder(ModelMixin, ConfigMixin):
+    """
+    Wrapper for CLIP sentence encoder using sentence-transformers.
+
+    Generates text embeddings for cross-attention conditioning.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/clip-ViT-B-16",
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.model = None  # Lazy loading
+
+    def _load_model(self):
+        """Lazy load the sentence transformer model."""
+        if self.model is None:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model_name)
+            self.model.eval()
+
+    def forward(self, text: Union[str, List[str]]) -> torch.Tensor:
+        """
+        Encode text using CLIP sentence transformer.
+
+        Args:
+            text: Input text or list of texts.
+
+        Returns:
+            Text embedding of shape (B, 512).
+        """
+        self._load_model()
+
+        if isinstance(text, str):
+            text = [text]
+
+        embeddings = self.model.encode(text, convert_to_tensor=True)
+        return embeddings
+
+
+class MultiConditionEncoder(ModelMixin, ConfigMixin):
+    """
+    Multi-condition encoder that combines all conditioning modalities.
+
+    This encoder takes multiple condition inputs and produces:
+    - c_crossattn: Features for cross-attention (image, text, palette embeddings)
+    - c_concat: Features for concatenation (sketch encoding)
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sketch_in_channels: int = 1,
+        sketch_out_channels: int = 4,
+        palette_in_channels: int = 3,
+        palette_hidden_channels: int = 64,
+        palette_out_channels: int = 512,
+        n_colors: int = 5,
+        clip_image_model: str = "ViT-B/16",
+        clip_text_model: str = "sentence-transformers/clip-ViT-B-16",
+    ):
+        super().__init__()
+
+        self.sketch_encoder = SketchEncoder(
+            in_channels=sketch_in_channels,
+            out_channels=sketch_out_channels,
+        )
+
+        self.palette_encoder = PaletteEncoder(
+            in_channels=palette_in_channels,
+            hidden_channels=palette_hidden_channels,
+            out_channels=palette_out_channels,
+            n_colors=n_colors,
+        )
+
+        # CLIP encoders are lazy-loaded
+        self.clip_image_encoder = None
+        self.clip_text_encoder = None
+        self._clip_image_model = clip_image_model
+        self._clip_text_model = clip_text_model
+
+    def _load_clip_encoders(self):
+        """Lazy load CLIP encoders."""
+        if self.clip_image_encoder is None:
+            self.clip_image_encoder = CLIPImageEncoder(
+                model_name=self._clip_image_model
+            )
+        if self.clip_text_encoder is None:
+            self.clip_text_encoder = CLIPTextEncoder(model_name=self._clip_text_model)
+
+    def encode_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Encode image using CLIP."""
+        self._load_clip_encoders()
+        return self.clip_image_encoder(image)
+
+    def encode_text(self, text: Union[str, List[str]]) -> torch.Tensor:
+        """Encode text using CLIP."""
+        self._load_clip_encoders()
+        return self.clip_text_encoder(text)
+
+    def encode_sketch(self, sketch: torch.Tensor) -> torch.Tensor:
+        """Encode sketch/edge map."""
+        return self.sketch_encoder(sketch)
+
+    def encode_palette(self, palette: torch.Tensor) -> torch.Tensor:
+        """Encode color palette."""
+        return self.palette_encoder(palette)
+
+    def get_unconditional_conditioning(
+        self,
+        batch_size: int = 1,
+        image_size: int = 256,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Get unconditional conditioning for classifier-free guidance.
+
+        IMPORTANT: The original model was trained to drop conditions by replacing them
+        with encoded placeholders (zero/gray image through CLIP, empty string through
+        sentence-transformers, zero palette through PaletteEncoder, zero sketch through
+        SketchEncoder) — NOT with zero tensors. This method produces the correct
+        unconditional embeddings.
+
+        Args:
+            batch_size: Batch size.
+            image_size: Image resolution (for sketch spatial dims).
+            device: Device to place tensors on.
+
+        Returns:
+            Dictionary with c_crossattn and c_concat for unconditional guidance.
+        """
+        return self.forward(
+            image_embed=None,
+            text=None,
+            sketch=None,
+            palette=None,
+            batch_size=batch_size,
+            image_size=image_size,
+            device=device,
+        )
+
+    def forward(
+        self,
+        image_embed: Optional[torch.Tensor] = None,
+        text: Optional[Union[str, List[str]]] = None,
+        sketch: Optional[torch.Tensor] = None,
+        palette: Optional[torch.Tensor] = None,
+        batch_size: int = 1,
+        image_size: int = 256,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Encode all conditions.
+
+        When a condition is not provided, the model encodes a placeholder input
+        through the actual encoder (matching training behavior) rather than using
+        zero tensors. This is critical because the model was trained with:
+        - Image drop → CLIP encoding of a gray/zero image (0.0 in [-1,1])
+        - Text drop → sentence-transformer encoding of ""
+        - Palette drop → PaletteEncoder(zeros)
+        - Sketch drop → SketchEncoder(zeros)
+
+        Args:
+            image_embed: Reference image of shape (B, 3, H, W) in [-1, 1].
+            text: Text description(s).
+            sketch: Binary sketch of shape (B, 1, H, W) in [0, 1].
+            palette: Color palette of shape (B, n_colors, 3) in [0, 1].
+            batch_size: Batch size (used when no inputs are provided).
+            image_size: Image resolution (used to create placeholder sketch).
+            device: Device to place tensors on.
+
+        Returns:
+            Dictionary with:
+            - c_crossattn: Cross-attention context of shape (B, 3, 512) - always 3 tokens.
+            - c_concat: Concatenation features of shape (B, 4, H/8, W/8).
+        """
+        self._load_clip_encoders()
+
+        # Determine batch size and device from any available input
+        if image_embed is not None:
+            batch_size = image_embed.shape[0]
+            device = device or image_embed.device
+            image_size = image_embed.shape[-1]
+        elif sketch is not None:
+            batch_size = sketch.shape[0]
+            device = device or sketch.device
+            image_size = sketch.shape[-1]
+        elif palette is not None:
+            batch_size = palette.shape[0]
+            device = device or palette.device
+
+        device = device or torch.device("cpu")
+        # Infer dtype from model weights for placeholder tensors (e.g. float16)
+        dtype = next(self.palette_encoder.parameters()).dtype
+
+        # --- Image embedding (token 0) ---
+        # When not provided, encode a zero (gray) image through CLIP, matching training ucg_training val=0.0
+        if image_embed is not None:
+            img_emb = self.clip_image_encoder(image_embed)  # (B, 1, 512)
+        else:
+            placeholder_img = torch.zeros(
+                batch_size, 3, image_size, image_size, device=device, dtype=dtype
+            )
+            img_emb = self.clip_image_encoder(placeholder_img)  # (B, 1, 512)
+
+        # --- Text embedding (token 1) ---
+        # When not provided, encode empty string through sentence-transformers, matching training ucg_training val=""
+        if text is not None:
+            text_emb = self.clip_text_encoder(text)  # (B, 512)
+            if device is not None:
+                text_emb = text_emb.to(device)
+            text_emb = text_emb.unsqueeze(1)  # (B, 1, 512)
+        else:
+            text_emb = self.clip_text_encoder([""] * batch_size)  # (B, 512)
+            text_emb = text_emb.to(device).unsqueeze(1)  # (B, 1, 512)
+
+        # --- Palette embedding (token 2) ---
+        # When not provided, encode zero palette through PaletteEncoder, matching training ucg_training val=0.0
+        if palette is not None:
+            palette_emb = self.palette_encoder(palette)  # (B, 512)
+            palette_emb = palette_emb.unsqueeze(1)  # (B, 1, 512)
+        else:
+            n_colors = self.config.get("n_colors", 5)
+            placeholder_palette = torch.zeros(batch_size, n_colors, 3, device=device, dtype=dtype)
+            palette_emb = self.palette_encoder(placeholder_palette)  # (B, 512)
+            palette_emb = palette_emb.unsqueeze(1)  # (B, 1, 512)
+
+        # Combine cross-attention embeddings - always (B, 3, 512)
+        c_crossattn = torch.cat([img_emb, text_emb, palette_emb], dim=1)
+
+        # --- Sketch encoding for concatenation ---
+        # When not provided, encode zero sketch through SketchEncoder, matching training ucg_training val=0.0
+        if sketch is not None:
+            c_concat = self.sketch_encoder(sketch)  # (B, 4, H/8, W/8)
+        else:
+            placeholder_sketch = torch.zeros(
+                batch_size, 1, image_size, image_size, device=device, dtype=dtype
+            )
+            c_concat = self.sketch_encoder(placeholder_sketch)  # (B, 4, H/8, W/8)
+
+        return {
+            "c_crossattn": c_crossattn,
+            "c_concat": c_concat,
+        }

model_index.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": [
+    "pipeline_matfuse",
+    "MatFusePipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "vae_matfuse",
+    "MatFuseVQModel"
+  ],
+  "condition_encoder": [
+    "condition_encoders",
+    "MultiConditionEncoder"
+  ],
+  "scheduler": [
+    "diffusers",
+    "DDIMScheduler"
+  ]
+}

pipeline_matfuse.py

@@ -0,0 +1,535 @@
+"""
+MatFuse Pipeline for diffusers.
+
+A custom diffusers pipeline for generating PBR material maps using the MatFuse model.
+
+Note: This pipeline uses:
+- Standard UNet2DConditionModel from diffusers (with custom in/out channels config)
+- Custom MatFuseVQModel (required because MatFuse uses 4 separate encoders/quantizers)
+"""
+
+import os
+import inspect
+from typing import Optional, Union, List, Callable, Dict, Any, Tuple
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import numpy as np
+
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.models import UNet2DConditionModel
+from diffusers.schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler
+
+try:
+    from vae_matfuse import MatFuseVQModel
+except ImportError:
+    from diffusers.models.modeling_utils import ModelMixin as MatFuseVQModel
+try:
+    from condition_encoders import MultiConditionEncoder
+except ImportError:
+    from diffusers.models.modeling_utils import ModelMixin as MultiConditionEncoder
+
+
+class MatFusePipeline(DiffusionPipeline):
+    """
+    Pipeline for generating PBR material maps using MatFuse.
+    
+    This pipeline generates 4 material maps (diffuse, normal, roughness, specular)
+    from various conditioning inputs like reference images, text, sketches, and color palettes.
+    
+    Args:
+        vae: MatFuseVQModel for encoding/decoding material maps (custom, required).
+        unet: UNet2DConditionModel for denoising (standard diffusers model).
+        scheduler: Diffusion scheduler.
+        condition_encoder: Multi-condition encoder for processing inputs.
+    
+    Note:
+        The VQ-VAE must be the custom MatFuseVQModel because MatFuse uses 4 separate
+        encoders and quantizers (one per material map type). The UNet can be the
+        standard diffusers UNet2DConditionModel configured with:
+        - in_channels=16 (12 latent + 4 sketch concat)
+        - out_channels=12 (4 maps × 3 channels)
+        - cross_attention_dim=512
+    """
+    
+    model_cpu_offload_seq = "condition_encoder->unet->vae"
+    _optional_components = ["condition_encoder"]
+    
+    def __init__(
+        self,
+        vae: MatFuseVQModel,
+        unet: UNet2DConditionModel,
+        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler],
+        condition_encoder: Optional[MultiConditionEncoder] = None,
+    ):
+        super().__init__()
+        
+        self.register_modules(
+            vae=vae,
+            unet=unet,
+            scheduler=scheduler,
+            condition_encoder=condition_encoder,
+        )
+        
+        self.vae_scale_factor = 8  # Downsampling factor of VQ-VAE
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Load the MatFuse pipeline from a local directory.
+        
+        Loads each component (UNet, VAE, scheduler, condition_encoder) individually
+        from their respective subdirectories.
+        
+        Args:
+            pretrained_model_name_or_path: Path to the directory containing the model components.
+            **kwargs: Additional keyword arguments (e.g., torch_dtype).
+        """
+        model_dir = pretrained_model_name_or_path
+        torch_dtype = kwargs.get("torch_dtype", None)
+        
+        # Load UNet (standard diffusers)
+        unet = UNet2DConditionModel.from_pretrained(
+            os.path.join(model_dir, "unet"),
+            torch_dtype=torch_dtype,
+        )
+        
+        # Load VAE (custom)
+        vae = MatFuseVQModel.from_pretrained(
+            os.path.join(model_dir, "vae"),
+            torch_dtype=torch_dtype,
+        )
+        
+        # Load scheduler
+        scheduler = DDIMScheduler.from_pretrained(
+            os.path.join(model_dir, "scheduler"),
+        )
+        
+        # Load condition encoder (custom) if it exists
+        cond_dir = os.path.join(model_dir, "condition_encoder")
+        condition_encoder = None
+        if os.path.isdir(cond_dir):
+            condition_encoder = MultiConditionEncoder.from_pretrained(
+                cond_dir,
+                torch_dtype=torch_dtype,
+            )
+        
+        return cls(
+            vae=vae,
+            unet=unet,
+            scheduler=scheduler,
+            condition_encoder=condition_encoder,
+        )
+        
+    @property
+    def _execution_device(self):
+        if self.device != torch.device("meta"):
+            return self.device
+        for name, model in self.components.items():
+            if isinstance(model, torch.nn.Module):
+                return next(model.parameters()).device
+        # Also check condition_encoder (may not be in components dict)
+        if self.condition_encoder is not None:
+            return next(self.condition_encoder.parameters()).device
+        return torch.device("cpu")
+
+    def to(self, *args, **kwargs):
+        """Override to() to also move condition_encoder (not auto-tracked by diffusers)."""
+        result = super().to(*args, **kwargs)
+        if self.condition_encoder is not None:
+            self.condition_encoder = self.condition_encoder.to(*args, **kwargs)
+        return result
+
+    def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode latents to material maps."""
+        # Add circular padding for seamless textures
+        latents = F.pad(latents, (7, 7, 7, 7), mode="circular")
+        
+        # Upcast to float32 for VAE decoding to avoid NaN from float16 precision
+        needs_upcast = latents.dtype == torch.float16
+        if needs_upcast:
+            self.vae.to(dtype=torch.float32)
+            latents = latents.float()
+        
+        # Decode
+        materials = self.vae.decode(latents)
+        
+        if needs_upcast:
+            self.vae.to(dtype=torch.float16)
+            materials = materials.half()
+        
+        # Center crop to remove padding
+        _, _, h, w = materials.shape
+        target_h = (h - 14 * self.vae_scale_factor)
+        target_w = (w - 14 * self.vae_scale_factor)
+        start_h = (h - target_h) // 2
+        start_w = (w - target_w) // 2
+        materials = materials[:, :, start_h:start_h + target_h, start_w:start_w + target_w]
+        
+        return materials
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Prepare initial noise latents."""
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        
+        if latents is None:
+            latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(device)
+        
+        # Scale by scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        
+        return latents
+
+    def prepare_extra_step_kwargs(self, generator: Optional[torch.Generator], eta: float) -> Dict[str, Any]:
+        """Prepare extra kwargs for the scheduler step."""
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        
+        return extra_step_kwargs
+
+    def _encode_conditions(
+        self,
+        image: Optional[torch.Tensor] = None,
+        text: Optional[Union[str, List[str]]] = None,
+        sketch: Optional[torch.Tensor] = None,
+        palette: Optional[torch.Tensor] = None,
+        batch_size: int = 1,
+        image_size: int = 256,
+        device: torch.device = None,
+        dtype: torch.dtype = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode all condition inputs through their respective encoders.
+        
+        When a condition is not provided, the encoder creates a placeholder
+        and encodes it (matching training behavior), rather than using zero tensors.
+        """
+        device = device or self._execution_device
+        
+        if self.condition_encoder is not None:
+            cond = self.condition_encoder(
+                image_embed=image,
+                text=text,
+                sketch=sketch,
+                palette=palette,
+                batch_size=batch_size,
+                image_size=image_size,
+                device=device,
+            )
+            c_crossattn = cond["c_crossattn"]
+            c_concat = cond["c_concat"]
+        else:
+            c_crossattn = None
+            c_concat = None
+        
+        # Ensure proper dtype
+        if c_crossattn is not None:
+            c_crossattn = c_crossattn.to(dtype=dtype, device=device)
+        if c_concat is not None:
+            c_concat = c_concat.to(dtype=dtype, device=device)
+        
+        return c_crossattn, c_concat
+
+    def _get_uncond_embeddings(
+        self,
+        batch_size: int,
+        image_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Get unconditional embeddings for classifier-free guidance.
+        
+        Creates proper unconditional embeddings by encoding placeholder inputs
+        through the actual encoders (gray image → CLIP, empty string → SentenceTransformer,
+        zero palette → PaletteEncoder, zero sketch → SketchEncoder).
+        
+        This matches the original training behavior where ucg_training drops conditions
+        by setting them to val=0.0 (images/palette/sketch) or val="" (text), and then
+        encoding those placeholder values through the encoders.
+        """
+        if self.condition_encoder is not None:
+            uc = self.condition_encoder.get_unconditional_conditioning(
+                batch_size=batch_size,
+                image_size=image_size,
+                device=device,
+            )
+            uc_crossattn = uc["c_crossattn"].to(dtype=dtype, device=device)
+            uc_concat = uc["c_concat"].to(dtype=dtype, device=device)
+        else:
+            uc_crossattn = None
+            uc_concat = None
+        
+        return uc_crossattn, uc_concat
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Optional[Union[torch.Tensor, Image.Image]] = None,
+        text: Optional[Union[str, List[str]]] = None,
+        sketch: Optional[Union[torch.Tensor, Image.Image]] = None,
+        palette: Optional[Union[torch.Tensor, np.ndarray, List[Tuple[int, int, int]]]] = None,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        num_images_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+        callback_steps: int = 1,
+    ) -> Dict[str, Any]:
+        """
+        Generate PBR material maps.
+        
+        Args:
+            image: Reference image for style/appearance guidance.
+            text: Text description of the material.
+            sketch: Binary edge/sketch map for structure guidance.
+            palette: Color palette (5 colors) for color guidance.
+            height: Output image height.
+            width: Output image width.
+            num_inference_steps: Number of denoising steps.
+            guidance_scale: Classifier-free guidance scale.
+            num_images_per_prompt: Number of images to generate per prompt.
+            eta: DDIM eta parameter.
+            generator: Random number generator for reproducibility.
+            latents: Pre-generated noise latents.
+            output_type: Output format ("pil", "tensor", "np").
+            return_dict: Whether to return a dict.
+            callback: Callback function called every `callback_steps` steps.
+            callback_steps: Frequency of callback calls.
+            
+        Returns:
+            Dictionary containing:
+            - images: List of generated images (4 maps per generation).
+            - diffuse: Diffuse/albedo maps.
+            - normal: Normal maps.
+            - roughness: Roughness maps.
+            - specular: Specular maps.
+        """
+        device = self._execution_device
+        dtype = self.unet.dtype if hasattr(self.unet, 'dtype') else torch.float32
+        
+        # Determine batch size
+        if text is not None and isinstance(text, str):
+            batch_size = 1
+        elif text is not None:
+            batch_size = len(text)
+        else:
+            batch_size = 1
+        
+        batch_size = batch_size * num_images_per_prompt
+        
+        # Preprocess inputs
+        if image is not None and isinstance(image, Image.Image):
+            image = self._preprocess_image(image, device, dtype)
+        
+        if sketch is not None and isinstance(sketch, Image.Image):
+            sketch = self._preprocess_sketch(sketch, height, width, device, dtype)
+        
+        if palette is not None and not isinstance(palette, torch.Tensor):
+            palette = self._preprocess_palette(palette, device, dtype)
+        
+        # Encode conditions
+        # The encoder handles None conditions by encoding placeholder inputs
+        # (matching the original model's UCG training behavior)
+        c_crossattn, c_concat = self._encode_conditions(
+            image=image,
+            text=text,
+            sketch=sketch,
+            palette=palette,
+            batch_size=batch_size,
+            image_size=height,
+            device=device,
+            dtype=dtype,
+        )
+        
+        # Get unconditional embeddings for CFG
+        # These are encoded placeholders, NOT zero tensors
+        do_classifier_free_guidance = guidance_scale > 1.0
+        if do_classifier_free_guidance:
+            uc_crossattn, uc_concat = self._get_uncond_embeddings(
+                batch_size, height, device, dtype
+            )
+        
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        
+        # Prepare latent variables
+        num_channels_latents = 12  # 4 maps * 3 channels per quantizer
+        latents = self.prepare_latents(
+            batch_size,
+            num_channels_latents,
+            height,
+            width,
+            dtype,
+            device,
+            generator,
+            latents,
+        )
+        
+        # Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        
+        # Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Prepare latent input with sketch conditioning
+                if do_classifier_free_guidance:
+                    # For CFG: unconditional uses uc_concat, conditional uses c_concat
+                    latent_uncond = torch.cat([latents, uc_concat], dim=1)
+                    latent_cond = torch.cat([latents, c_concat], dim=1)
+                    latent_model_input = torch.cat([latent_uncond, latent_cond])
+                    if c_crossattn is not None:
+                        encoder_hidden_states = torch.cat([uc_crossattn, c_crossattn])
+                    else:
+                        encoder_hidden_states = None
+                else:
+                    latent_model_input = torch.cat([latents, c_concat], dim=1)
+                    encoder_hidden_states = c_crossattn
+                
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                
+                # Predict noise
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=encoder_hidden_states,
+                    return_dict=False,
+                )
+                # return_dict=False returns tuple, first element is sample
+                if isinstance(noise_pred, tuple):
+                    noise_pred = noise_pred[0]
+                elif isinstance(noise_pred, dict):
+                    noise_pred = noise_pred["sample"]
+                
+                # Classifier-free guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                
+                # Compute previous noisy sample
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                
+                # Callback
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+        
+        # Decode latents
+        materials = self.decode_latents(latents)
+        
+        # Split into individual maps
+        diffuse = materials[:, 0:3]
+        normal = materials[:, 3:6]
+        roughness = materials[:, 6:9]
+        specular = materials[:, 9:12]
+        
+        # Post-process outputs
+        if output_type == "pil":
+            diffuse = self._tensor_to_pil(diffuse)
+            normal = self._tensor_to_pil(normal)
+            roughness = self._tensor_to_pil(roughness)
+            specular = self._tensor_to_pil(specular)
+        elif output_type == "np":
+            diffuse = self._tensor_to_numpy(diffuse)
+            normal = self._tensor_to_numpy(normal)
+            roughness = self._tensor_to_numpy(roughness)
+            specular = self._tensor_to_numpy(specular)
+        
+        if return_dict:
+            return {
+                "diffuse": diffuse,
+                "normal": normal,
+                "roughness": roughness,
+                "specular": specular,
+            }
+        
+        return (diffuse, normal, roughness, specular)
+
+    def _preprocess_image(self, image: Image.Image, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        """Preprocess PIL image to tensor."""
+        image = image.convert("RGB")
+        image = np.array(image).astype(np.float32) / 255.0
+        image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0)
+        image = image * 2.0 - 1.0  # Scale to [-1, 1]
+        return image.to(device=device, dtype=dtype)
+
+    def _preprocess_sketch(
+        self,
+        sketch: Image.Image,
+        height: int,
+        width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Preprocess sketch image to tensor."""
+        sketch = sketch.convert("L")
+        sketch = sketch.resize((width, height), Image.BILINEAR)
+        sketch = np.array(sketch).astype(np.float32) / 255.0
+        sketch = torch.from_numpy(sketch).unsqueeze(0).unsqueeze(0)
+        return sketch.to(device=device, dtype=dtype)
+
+    def _preprocess_palette(
+        self,
+        palette: Union[np.ndarray, List[Tuple[int, int, int]]],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Preprocess color palette to tensor."""
+        if isinstance(palette, list):
+            palette = np.array(palette, dtype=np.float32) / 255.0
+        elif isinstance(palette, np.ndarray):
+            if palette.max() > 1.0:
+                palette = palette.astype(np.float32) / 255.0
+            else:
+                palette = palette.astype(np.float32)
+        
+        # Ensure 5 colors
+        while len(palette) < 5:
+            palette = np.concatenate([palette, palette[-1:]], axis=0)
+        palette = palette[:5]
+        
+        palette = torch.from_numpy(palette).unsqueeze(0)
+        return palette.to(device=device, dtype=dtype)
+
+    def _tensor_to_pil(self, tensor: torch.Tensor) -> List[Image.Image]:
+        """Convert tensor to list of PIL images."""
+        tensor = (tensor + 1.0) / 2.0
+        tensor = tensor.clamp(0, 1)
+        tensor = tensor.cpu().permute(0, 2, 3, 1).numpy()
+        tensor = (tensor * 255).astype(np.uint8)
+        return [Image.fromarray(img) for img in tensor]
+
+    def _tensor_to_numpy(self, tensor: torch.Tensor) -> np.ndarray:
+        """Convert tensor to numpy array."""
+        tensor = (tensor + 1.0) / 2.0
+        tensor = tensor.clamp(0, 1)
+        return tensor.cpu().permute(0, 2, 3, 1).numpy()

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.35.2",
+  "beta_end": 0.0195,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.0015,
+  "clip_sample": false,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "set_alpha_to_one": false,
+  "steps_offset": 0,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}

unet/config.json

@@ -0,0 +1,64 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.35.2",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": [8, 16, 32],
+  "attention_type": "default",
+  "block_out_channels": [
+    256,
+    512,
+    1024
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 512,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 16,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 12,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 32,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

unet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de0fd0026e9621a105c46aea51a671dae90b5e1233ddf0c4e26a54eb0097e3d1
+size 1580197816

vae/config.json

@@ -0,0 +1,21 @@
+{
+  "_class_name": "MatFuseVQModel",
+  "_diffusers_version": "0.35.2",
+  "attn_resolutions": [],
+  "ch": 128,
+  "ch_mult": [
+    1,
+    1,
+    2,
+    4
+  ],
+  "dropout": 0.0,
+  "embed_dim": 3,
+  "in_channels": 3,
+  "n_embed": 4096,
+  "num_res_blocks": 2,
+  "out_channels": 12,
+  "resolution": 256,
+  "scaling_factor": 1.0,
+  "z_channels": 256
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad0962b5d9ad812c2ee8fe146e28ee73bba4382be388cc6703e3b4bb6ce6d06b
+size 528848272

vae/vae_matfuse.py

@@ -0,0 +1,590 @@
+"""
+MatFuse VQ-VAE Model for diffusers.
+
+This is a custom VQ-VAE that has 4 separate encoders (one for each material map)
+and 4 separate quantizers, with a single shared decoder.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def Normalize(in_channels: int, num_groups: int = 32) -> nn.GroupNorm:
+    """Group normalization."""
+    return nn.GroupNorm(
+        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    """Swish activation."""
+    return x * torch.sigmoid(x)
+
+
+class Upsample(nn.Module):
+    """Upsampling layer with optional convolution."""
+
+    def __init__(self, in_channels: int, with_conv: bool = True):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """Downsampling layer with optional convolution."""
+
+    def __init__(self, in_channels: int, with_conv: bool = True):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    """Residual block with optional time embedding."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(
+        self, x: torch.Tensor, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None and hasattr(self, "temb_proj"):
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Self-attention block."""
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # Compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)  # b, hw, c
+        k = k.reshape(b, c, h * w)  # b, c, hw
+        w_ = torch.bmm(q, k)  # b, hw, hw
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # Attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b, hw, hw
+        h_ = torch.bmm(v, w_)  # b, c, hw
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class Encoder(nn.Module):
+    """Encoder module for VQ-VAE."""
+
+    def __init__(
+        self,
+        ch: int = 128,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        resolution: int = 256,
+        z_channels: int = 256,
+        double_z: bool = False,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # Downsampling
+        self.conv_in = nn.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, with_conv=True)
+                curr_res = curr_res // 2
+
+            self.down.append(down)
+
+        # Middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # End
+        self.norm_out = Normalize(block_in)
+        out_channels = 2 * z_channels if double_z else z_channels
+        self.conv_out = nn.Conv2d(
+            block_in, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Downsampling
+        h = self.conv_in(x)
+
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h, None)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+            if hasattr(self.down[i_level], "downsample"):
+                h = self.down[i_level].downsample(h)
+
+        # Middle
+        h = self.mid.block_1(h, None)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, None)
+
+        # End
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class Decoder(nn.Module):
+    """Decoder module for VQ-VAE."""
+
+    def __init__(
+        self,
+        ch: int = 128,
+        out_ch: int = 12,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        resolution: int = 256,
+        z_channels: int = 256,
+        give_pre_end: bool = False,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+
+        # Compute in_ch_mult and block_in
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // (2 ** (self.num_resolutions - 1))
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # Middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # Upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+
+            if i_level != 0:
+                up.upsample = Upsample(block_in, with_conv=True)
+                curr_res = curr_res * 2
+
+            self.up.insert(0, up)
+
+        # End
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # Middle
+        h = self.mid.block_1(h, None)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, None)
+
+        # Upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, None)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+
+        # End
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class VectorQuantizer(nn.Module):
+    """
+    Vector Quantizer module.
+
+    Discretizes the input vectors using a learned codebook.
+    """
+
+    def __init__(
+        self,
+        n_embed: int,
+        embed_dim: int,
+        beta: float = 0.25,
+    ):
+        super().__init__()
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.beta = beta
+
+        self.embedding = nn.Embedding(self.n_embed, self.embed_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_embed, 1.0 / self.n_embed)
+
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
+        # Reshape z -> (batch, height, width, channel) and flatten
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.embed_dim)
+
+        # Distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+        d = (
+            torch.sum(z_flattened**2, dim=1, keepdim=True)
+            + torch.sum(self.embedding.weight**2, dim=1)
+            - 2 * torch.einsum("bd,dn->bn", z_flattened, self.embedding.weight.t())
+        )
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+
+        # Compute loss for embedding
+        loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean(
+            (z_q - z.detach()) ** 2
+        )
+
+        # Preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # Reshape back to match original input shape
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q, loss, (None, None, min_encoding_indices)
+
+    def get_codebook_entry(
+        self, indices: torch.Tensor, shape: Optional[Tuple] = None
+    ) -> torch.Tensor:
+        # Get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+            # Reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+class MatFuseVQModel(ModelMixin, ConfigMixin):
+    """
+    MatFuse VQ-VAE Model.
+
+    This model has 4 separate encoders for each material map (diffuse, normal, roughness, specular)
+    and 4 separate VQ quantizers, with a single shared decoder that outputs 12 channels.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        ch: int = 128,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        out_channels: int = 12,
+        resolution: int = 256,
+        z_channels: int = 256,
+        n_embed: int = 4096,
+        embed_dim: int = 3,
+        scaling_factor: float = 1.0,
+    ):
+        super().__init__()
+
+        self.scaling_factor = scaling_factor
+        self.embed_dim = embed_dim
+
+        ddconfig = dict(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolutions,
+            dropout=dropout,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+            double_z=False,
+        )
+
+        # 4 separate encoders for each material map
+        self.encoder_0 = Encoder(**ddconfig)
+        self.encoder_1 = Encoder(**ddconfig)
+        self.encoder_2 = Encoder(**ddconfig)
+        self.encoder_3 = Encoder(**ddconfig)
+
+        # Single decoder
+        decoder_config = dict(
+            ch=ch,
+            out_ch=out_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolutions,
+            dropout=dropout,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(**decoder_config)
+
+        # 4 separate quantizers
+        self.quantize_0 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_1 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_2 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_3 = VectorQuantizer(n_embed, embed_dim)
+
+        # Quant convolutions
+        self.quant_conv_0 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_1 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_2 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_3 = nn.Conv2d(z_channels, embed_dim, 1)
+
+        # Post quant convolution (takes 4 * embed_dim channels)
+        self.post_quant_conv = nn.Conv2d(embed_dim * 4, z_channels, 1)
+
+    def encode_to_prequant(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to pre-quantized latent space."""
+        h_0 = self.encoder_0(x[:, :3])
+        h_1 = self.encoder_1(x[:, 3:6])
+        h_2 = self.encoder_2(x[:, 6:9])
+        h_3 = self.encoder_3(x[:, 9:12])
+
+        h_0 = self.quant_conv_0(h_0)
+        h_1 = self.quant_conv_1(h_1)
+        h_2 = self.quant_conv_2(h_2)
+        h_3 = self.quant_conv_3(h_3)
+
+        h = torch.cat((h_0, h_1, h_2, h_3), dim=1)
+        return h
+
+    def quantize_latent(
+        self, h: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Quantize the latent space."""
+        quant_0, emb_loss_0, info_0 = self.quantize_0(h[:, : self.embed_dim])
+        quant_1, emb_loss_1, info_1 = self.quantize_1(
+            h[:, self.embed_dim : 2 * self.embed_dim]
+        )
+        quant_2, emb_loss_2, info_2 = self.quantize_2(
+            h[:, 2 * self.embed_dim : 3 * self.embed_dim]
+        )
+        quant_3, emb_loss_3, info_3 = self.quantize_3(h[:, 3 * self.embed_dim :])
+
+        quant = torch.cat((quant_0, quant_1, quant_2, quant_3), dim=1)
+        emb_loss = emb_loss_0 + emb_loss_1 + emb_loss_2 + emb_loss_3
+        info = torch.stack([info_0[-1], info_1[-1], info_2[-1], info_3[-1]], dim=0)
+
+        return quant, emb_loss, info
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to quantized latent space."""
+        h = self.encode_to_prequant(x)
+        quant, _, _ = self.quantize_latent(h)
+        return quant * self.scaling_factor
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode from latent space to image."""
+        z = z / self.scaling_factor
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass through the VQ-VAE."""
+        h = self.encode_to_prequant(x)
+        quant, diff, _ = self.quantize_latent(h)
+        dec = self.decode(quant * self.scaling_factor)
+        return dec, diff

vae_matfuse.py

@@ -0,0 +1,590 @@
+"""
+MatFuse VQ-VAE Model for diffusers.
+
+This is a custom VQ-VAE that has 4 separate encoders (one for each material map)
+and 4 separate quantizers, with a single shared decoder.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def Normalize(in_channels: int, num_groups: int = 32) -> nn.GroupNorm:
+    """Group normalization."""
+    return nn.GroupNorm(
+        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    """Swish activation."""
+    return x * torch.sigmoid(x)
+
+
+class Upsample(nn.Module):
+    """Upsampling layer with optional convolution."""
+
+    def __init__(self, in_channels: int, with_conv: bool = True):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """Downsampling layer with optional convolution."""
+
+    def __init__(self, in_channels: int, with_conv: bool = True):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    """Residual block with optional time embedding."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(
+        self, x: torch.Tensor, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None and hasattr(self, "temb_proj"):
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Self-attention block."""
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # Compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)  # b, hw, c
+        k = k.reshape(b, c, h * w)  # b, c, hw
+        w_ = torch.bmm(q, k)  # b, hw, hw
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # Attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b, hw, hw
+        h_ = torch.bmm(v, w_)  # b, c, hw
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class Encoder(nn.Module):
+    """Encoder module for VQ-VAE."""
+
+    def __init__(
+        self,
+        ch: int = 128,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        resolution: int = 256,
+        z_channels: int = 256,
+        double_z: bool = False,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # Downsampling
+        self.conv_in = nn.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, with_conv=True)
+                curr_res = curr_res // 2
+
+            self.down.append(down)
+
+        # Middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # End
+        self.norm_out = Normalize(block_in)
+        out_channels = 2 * z_channels if double_z else z_channels
+        self.conv_out = nn.Conv2d(
+            block_in, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Downsampling
+        h = self.conv_in(x)
+
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h, None)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+            if hasattr(self.down[i_level], "downsample"):
+                h = self.down[i_level].downsample(h)
+
+        # Middle
+        h = self.mid.block_1(h, None)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, None)
+
+        # End
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class Decoder(nn.Module):
+    """Decoder module for VQ-VAE."""
+
+    def __init__(
+        self,
+        ch: int = 128,
+        out_ch: int = 12,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        resolution: int = 256,
+        z_channels: int = 256,
+        give_pre_end: bool = False,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+
+        # Compute in_ch_mult and block_in
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // (2 ** (self.num_resolutions - 1))
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # Middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # Upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+
+            if i_level != 0:
+                up.upsample = Upsample(block_in, with_conv=True)
+                curr_res = curr_res * 2
+
+            self.up.insert(0, up)
+
+        # End
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # Middle
+        h = self.mid.block_1(h, None)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, None)
+
+        # Upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, None)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+
+        # End
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class VectorQuantizer(nn.Module):
+    """
+    Vector Quantizer module.
+
+    Discretizes the input vectors using a learned codebook.
+    """
+
+    def __init__(
+        self,
+        n_embed: int,
+        embed_dim: int,
+        beta: float = 0.25,
+    ):
+        super().__init__()
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.beta = beta
+
+        self.embedding = nn.Embedding(self.n_embed, self.embed_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_embed, 1.0 / self.n_embed)
+
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]:
+        # Reshape z -> (batch, height, width, channel) and flatten
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.embed_dim)
+
+        # Distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+        d = (
+            torch.sum(z_flattened**2, dim=1, keepdim=True)
+            + torch.sum(self.embedding.weight**2, dim=1)
+            - 2 * torch.einsum("bd,dn->bn", z_flattened, self.embedding.weight.t())
+        )
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+
+        # Compute loss for embedding
+        loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean(
+            (z_q - z.detach()) ** 2
+        )
+
+        # Preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # Reshape back to match original input shape
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q, loss, (None, None, min_encoding_indices)
+
+    def get_codebook_entry(
+        self, indices: torch.Tensor, shape: Optional[Tuple] = None
+    ) -> torch.Tensor:
+        # Get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+            # Reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+class MatFuseVQModel(ModelMixin, ConfigMixin):
+    """
+    MatFuse VQ-VAE Model.
+
+    This model has 4 separate encoders for each material map (diffuse, normal, roughness, specular)
+    and 4 separate VQ quantizers, with a single shared decoder that outputs 12 channels.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        ch: int = 128,
+        ch_mult: Tuple[int, ...] = (1, 1, 2, 4),
+        num_res_blocks: int = 2,
+        attn_resolutions: Tuple[int, ...] = (),
+        dropout: float = 0.0,
+        in_channels: int = 3,
+        out_channels: int = 12,
+        resolution: int = 256,
+        z_channels: int = 256,
+        n_embed: int = 4096,
+        embed_dim: int = 3,
+        scaling_factor: float = 1.0,
+    ):
+        super().__init__()
+
+        self.scaling_factor = scaling_factor
+        self.embed_dim = embed_dim
+
+        ddconfig = dict(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolutions,
+            dropout=dropout,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+            double_z=False,
+        )
+
+        # 4 separate encoders for each material map
+        self.encoder_0 = Encoder(**ddconfig)
+        self.encoder_1 = Encoder(**ddconfig)
+        self.encoder_2 = Encoder(**ddconfig)
+        self.encoder_3 = Encoder(**ddconfig)
+
+        # Single decoder
+        decoder_config = dict(
+            ch=ch,
+            out_ch=out_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolutions,
+            dropout=dropout,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(**decoder_config)
+
+        # 4 separate quantizers
+        self.quantize_0 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_1 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_2 = VectorQuantizer(n_embed, embed_dim)
+        self.quantize_3 = VectorQuantizer(n_embed, embed_dim)
+
+        # Quant convolutions
+        self.quant_conv_0 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_1 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_2 = nn.Conv2d(z_channels, embed_dim, 1)
+        self.quant_conv_3 = nn.Conv2d(z_channels, embed_dim, 1)
+
+        # Post quant convolution (takes 4 * embed_dim channels)
+        self.post_quant_conv = nn.Conv2d(embed_dim * 4, z_channels, 1)
+
+    def encode_to_prequant(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to pre-quantized latent space."""
+        h_0 = self.encoder_0(x[:, :3])
+        h_1 = self.encoder_1(x[:, 3:6])
+        h_2 = self.encoder_2(x[:, 6:9])
+        h_3 = self.encoder_3(x[:, 9:12])
+
+        h_0 = self.quant_conv_0(h_0)
+        h_1 = self.quant_conv_1(h_1)
+        h_2 = self.quant_conv_2(h_2)
+        h_3 = self.quant_conv_3(h_3)
+
+        h = torch.cat((h_0, h_1, h_2, h_3), dim=1)
+        return h
+
+    def quantize_latent(
+        self, h: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Quantize the latent space."""
+        quant_0, emb_loss_0, info_0 = self.quantize_0(h[:, : self.embed_dim])
+        quant_1, emb_loss_1, info_1 = self.quantize_1(
+            h[:, self.embed_dim : 2 * self.embed_dim]
+        )
+        quant_2, emb_loss_2, info_2 = self.quantize_2(
+            h[:, 2 * self.embed_dim : 3 * self.embed_dim]
+        )
+        quant_3, emb_loss_3, info_3 = self.quantize_3(h[:, 3 * self.embed_dim :])
+
+        quant = torch.cat((quant_0, quant_1, quant_2, quant_3), dim=1)
+        emb_loss = emb_loss_0 + emb_loss_1 + emb_loss_2 + emb_loss_3
+        info = torch.stack([info_0[-1], info_1[-1], info_2[-1], info_3[-1]], dim=0)
+
+        return quant, emb_loss, info
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to quantized latent space."""
+        h = self.encode_to_prequant(x)
+        quant, _, _ = self.quantize_latent(h)
+        return quant * self.scaling_factor
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode from latent space to image."""
+        z = z / self.scaling_factor
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass through the VQ-VAE."""
+        h = self.encode_to_prequant(x)
+        quant, diff, _ = self.quantize_latent(h)
+        dec = self.decode(quant * self.scaling_factor)
+        return dec, diff