Add files using upload-large-folder tool

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 donkeymouse
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LaVi-Bridge/modules/__init__.py

@@ -0,0 +1 @@
+"""LaVi-Bridge modules used by TherA inference."""

LaVi-Bridge/modules/adapters.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+from inspect import isfunction
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.utils import BaseOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+
+def default(val, d):
+    if val is not None: return val
+    return d() if isfunction(d) else d
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out, mult=4, dropout=0.1):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = GEGLU(dim, inner_dim)
+        self.net = nn.Sequential(
+            project_in, 
+            nn.Dropout(dropout), 
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+    
+
+@dataclass
+class TextAdapterOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class TextAdapter(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self, in_dim, int_dim, out_dim):
+        super().__init__()
+        self.in_dim = in_dim
+        self.ff1 = FeedForward(in_dim, int_dim)
+        self.ff2 = FeedForward(int_dim, out_dim)
+        self.norm1 = nn.LayerNorm(in_dim)
+        self.norm2 = nn.LayerNorm(int_dim)
+
+    def forward(self, x):
+        x = self.ff1(self.norm1(x))
+        x = self.ff2(self.norm2(x))
+        return TextAdapterOutput(x)

LaVi-Bridge/modules/lora.py

@@ -0,0 +1,1106 @@
+import json
+from itertools import groupby
+from typing import Dict, List, Optional, Set, Tuple, Type, Union
+
+import torch
+import torch.nn as nn
+
+try:
+    from safetensors.torch import safe_open
+    from safetensors.torch import save_file as safe_save
+
+    safetensors_available = True
+except ImportError:
+    from .safe_open import safe_open
+
+    def safe_save(
+        tensors: Dict[str, torch.Tensor],
+        filename: str,
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        raise EnvironmentError(
+            "Saving safetensors requires the safetensors library. Please install with pip or similar."
+        )
+
+    safetensors_available = False
+
+
+class LoraInjectedLinear(nn.Module):
+    def __init__(
+        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
+    ):
+        super().__init__()
+
+        if r > min(in_features, out_features):
+            raise ValueError(
+                f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
+            )
+        self.r = r
+        self.linear = nn.Linear(in_features, out_features, bias)
+        self.lora_down = nn.Linear(in_features, r, bias=False)
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Linear(r, out_features, bias=False)
+        self.scale = scale
+        self.selector = nn.Identity()
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.linear(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Linear(self.r, self.r, bias=False)
+        self.selector.weight.data = torch.diag(diag)
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+
+class LoraInjectedConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups: int = 1,
+        bias: bool = True,
+        r: int = 4,
+        dropout_p: float = 0.1,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        if r > min(in_channels, out_channels):
+            raise ValueError(
+                f"LoRA rank {r} must be less or equal than {min(in_channels, out_channels)}"
+            )
+        self.r = r
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+        self.lora_down = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=r,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=False,
+        )
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Conv2d(
+            in_channels=r,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector = nn.Identity()
+        self.scale = scale
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.conv(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Conv2d(
+            in_channels=self.r,
+            out_channels=self.r,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector.weight.data = torch.diag(diag)
+
+        # same device + dtype as lora_up
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+
+UNET_DEFAULT_TARGET_REPLACE = {"CrossAttention", "Attention", "GEGLU"}
+
+UNET_EXTENDED_TARGET_REPLACE = {"ResnetBlock2D", "CrossAttention", "Attention", "GEGLU"}
+
+TEXT_ENCODER_DEFAULT_TARGET_REPLACE = {"CLIPAttention"}
+
+TEXT_ENCODER_EXTENDED_TARGET_REPLACE = {"CLIPAttention"}
+
+DEFAULT_TARGET_REPLACE = UNET_DEFAULT_TARGET_REPLACE
+
+EMBED_FLAG = "<embed>"
+
+
+def _find_children(
+    model,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+):
+    """
+    Find all modules of a certain class (or union of classes).
+
+    Returns all matching modules, along with the parent of those moduless and the
+    names they are referenced by.
+    """
+    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
+    for parent in model.modules():
+        for name, module in parent.named_children():
+            if any([isinstance(module, _class) for _class in search_class]):
+                yield parent, name, module
+
+
+def _find_modules_v2(
+    model,
+    ancestor_class: Optional[Set[str]] = None,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+    exclude_children_of: Optional[List[Type[nn.Module]]] = [
+        LoraInjectedLinear,
+        LoraInjectedConv2d,
+    ],
+):
+    """
+    Find all modules of a certain class (or union of classes) that are direct or
+    indirect descendants of other modules of a certain class (or union of classes).
+
+    Returns all matching modules, along with the parent of those moduless and the
+    names they are referenced by.
+    """
+
+    # Get the targets we should replace all linears under
+    if ancestor_class is not None:
+        ancestors = (
+            module
+            for module in model.modules()
+            if module.__class__.__name__ in ancestor_class
+        )
+    else:
+        # this, incase you want to naively iterate over all modules.
+        ancestors = [module for module in model.modules()]
+
+    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
+    for ancestor in ancestors:
+        for fullname, module in ancestor.named_modules():
+            if any([isinstance(module, _class) for _class in search_class]):
+                # Find the direct parent if this is a descendant, not a child, of target
+                *path, name = fullname.split(".")
+                parent = ancestor
+                while path:
+                    parent = parent.get_submodule(path.pop(0))
+                # Skip this linear if it's a child of a LoraInjectedLinear
+                if exclude_children_of and any(
+                    [isinstance(parent, _class) for _class in exclude_children_of]
+                ):
+                    continue
+                # Otherwise, yield it
+                yield parent, name, module
+
+
+def _find_modules_old(
+    model,
+    ancestor_class: Set[str] = DEFAULT_TARGET_REPLACE,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+    exclude_children_of: Optional[List[Type[nn.Module]]] = [LoraInjectedLinear],
+):
+    ret = []
+    for _module in model.modules():
+        if _module.__class__.__name__ in ancestor_class:
+
+            for name, _child_module in _module.named_modules():
+                if _child_module.__class__ in search_class:
+                    ret.append((_module, name, _child_module))
+    print(ret)
+    return ret
+
+
+_find_modules = _find_modules_v2
+
+
+def inject_trainable_lora(
+    model: nn.Module,
+    target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
+    r: int = 4,
+    loras=None,  # path to lora .pt
+    verbose: bool = False,
+    dropout_p: float = 0.0,
+    scale: float = 1.0,
+):
+    """
+    inject lora into model, and returns lora parameter groups.
+    """
+
+    require_grad_params = []
+    names = []
+
+    if loras != None:
+        loras = torch.load(loras)
+
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[nn.Linear]
+    ):
+        weight = _child_module.weight
+        bias = _child_module.bias
+        if verbose:
+            print("LoRA Injection : injecting lora into ", name)
+            print("LoRA Injection : weight shape", weight.shape)
+        _tmp = LoraInjectedLinear(
+            _child_module.in_features,
+            _child_module.out_features,
+            _child_module.bias is not None,
+            r=r,
+            dropout_p=dropout_p,
+            scale=scale,
+        )
+        _tmp.linear.weight = weight
+        if bias is not None:
+            _tmp.linear.bias = bias
+
+        # switch the module
+        _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
+        _module._modules[name] = _tmp
+
+        require_grad_params.append(_module._modules[name].lora_up.parameters())
+        require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+        if loras != None:
+            _module._modules[name].lora_up.weight = loras.pop(0)
+            _module._modules[name].lora_down.weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight.requires_grad = True
+        _module._modules[name].lora_down.weight.requires_grad = True
+        names.append(name)
+
+    return require_grad_params, names
+
+
+def inject_trainable_lora_extended(
+    model: nn.Module,
+    target_replace_module: Set[str] = UNET_EXTENDED_TARGET_REPLACE,
+    r: int = 4,
+    loras=None,  # path to lora .pt
+):
+    """
+    inject lora into model, and returns lora parameter groups.
+    """
+
+    require_grad_params = []
+    names = []
+
+    if loras != None:
+        loras = torch.load(loras)
+
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[nn.Linear, nn.Conv2d]
+    ):
+        if _child_module.__class__ == nn.Linear:
+            weight = _child_module.weight
+            bias = _child_module.bias
+            _tmp = LoraInjectedLinear(
+                _child_module.in_features,
+                _child_module.out_features,
+                _child_module.bias is not None,
+                r=r,
+            )
+            _tmp.linear.weight = weight
+            if bias is not None:
+                _tmp.linear.bias = bias
+        elif _child_module.__class__ == nn.Conv2d:
+            weight = _child_module.weight
+            bias = _child_module.bias
+            _tmp = LoraInjectedConv2d(
+                _child_module.in_channels,
+                _child_module.out_channels,
+                _child_module.kernel_size,
+                _child_module.stride,
+                _child_module.padding,
+                _child_module.dilation,
+                _child_module.groups,
+                _child_module.bias is not None,
+                r=r,
+            )
+
+            _tmp.conv.weight = weight
+            if bias is not None:
+                _tmp.conv.bias = bias
+
+        # switch the module
+        _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
+        if bias is not None:
+            _tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
+
+        _module._modules[name] = _tmp
+
+        require_grad_params.append(_module._modules[name].lora_up.parameters())
+        require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+        if loras != None:
+            _module._modules[name].lora_up.weight = loras.pop(0)
+            _module._modules[name].lora_down.weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight.requires_grad = True
+        _module._modules[name].lora_down.weight.requires_grad = True
+        names.append(name)
+
+    return require_grad_params, names
+
+
+def extract_lora_ups_down(model, target_replace_module=DEFAULT_TARGET_REPLACE):
+
+    loras = []
+
+    for _m, _n, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
+    ):
+        loras.append((_child_module.lora_up, _child_module.lora_down))
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+
+def extract_lora_as_tensor(
+    model, target_replace_module=DEFAULT_TARGET_REPLACE, as_fp16=True
+):
+
+    loras = []
+
+    for _m, _n, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
+    ):
+        up, down = _child_module.realize_as_lora()
+        if as_fp16:
+            up = up.to(torch.float16)
+            down = down.to(torch.float16)
+
+        loras.append((up, down))
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+
+def save_lora_weight(
+    model,
+    path="./lora.pt",
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+):
+    weights = []
+    for _up, _down in extract_lora_ups_down(
+        model, target_replace_module=target_replace_module
+    ):
+        weights.append(_up.weight.to("cpu").to(torch.float16))
+        weights.append(_down.weight.to("cpu").to(torch.float16))
+
+    torch.save(weights, path)
+
+
+def save_lora_as_json(model, path="./lora.json"):
+    weights = []
+    for _up, _down in extract_lora_ups_down(model):
+        weights.append(_up.weight.detach().cpu().numpy().tolist())
+        weights.append(_down.weight.detach().cpu().numpy().tolist())
+
+    import json
+
+    with open(path, "w") as f:
+        json.dump(weights, f)
+
+
+def save_safeloras_with_embeds(
+    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
+    embeds: Dict[str, torch.Tensor] = {},
+    outpath="./lora.safetensors",
+):
+    """
+    Saves the Lora from multiple modules in a single safetensor file.
+
+    modelmap is a dictionary of {
+        "module name": (module, target_replace_module)
+    }
+    """
+    weights = {}
+    metadata = {}
+
+    for name, (model, target_replace_module) in modelmap.items():
+        metadata[name] = json.dumps(list(target_replace_module))
+
+        for i, (_up, _down) in enumerate(
+            extract_lora_as_tensor(model, target_replace_module)
+        ):
+            rank = _down.shape[0]
+
+            metadata[f"{name}:{i}:rank"] = str(rank)
+            weights[f"{name}:{i}:up"] = _up
+            weights[f"{name}:{i}:down"] = _down
+
+    for token, tensor in embeds.items():
+        metadata[token] = EMBED_FLAG
+        weights[token] = tensor
+
+    print(f"Saving weights to {outpath}")
+    safe_save(weights, outpath, metadata)
+
+
+def save_safeloras(
+    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
+    outpath="./lora.safetensors",
+):
+    return save_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
+
+
+def convert_loras_to_safeloras_with_embeds(
+    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
+    embeds: Dict[str, torch.Tensor] = {},
+    outpath="./lora.safetensors",
+):
+    """
+    Converts the Lora from multiple pytorch .pt files into a single safetensor file.
+
+    modelmap is a dictionary of {
+        "module name": (pytorch_model_path, target_replace_module, rank)
+    }
+    """
+
+    weights = {}
+    metadata = {}
+
+    for name, (path, target_replace_module, r) in modelmap.items():
+        metadata[name] = json.dumps(list(target_replace_module))
+
+        lora = torch.load(path)
+        for i, weight in enumerate(lora):
+            is_up = i % 2 == 0
+            i = i // 2
+
+            if is_up:
+                metadata[f"{name}:{i}:rank"] = str(r)
+                weights[f"{name}:{i}:up"] = weight
+            else:
+                weights[f"{name}:{i}:down"] = weight
+
+    for token, tensor in embeds.items():
+        metadata[token] = EMBED_FLAG
+        weights[token] = tensor
+
+    print(f"Saving weights to {outpath}")
+    safe_save(weights, outpath, metadata)
+
+
+def convert_loras_to_safeloras(
+    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
+    outpath="./lora.safetensors",
+):
+    convert_loras_to_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
+
+
+def parse_safeloras(
+    safeloras,
+) -> Dict[str, Tuple[List[nn.parameter.Parameter], List[int], List[str]]]:
+    """
+    Converts a loaded safetensor file that contains a set of module Loras
+    into Parameters and other information
+
+    Output is a dictionary of {
+        "module name": (
+            [list of weights],
+            [list of ranks],
+            target_replacement_modules
+        )
+    }
+    """
+    loras = {}
+    metadata = safeloras.metadata()
+
+    get_name = lambda k: k.split(":")[0]
+
+    keys = list(safeloras.keys())
+    keys.sort(key=get_name)
+
+    for name, module_keys in groupby(keys, get_name):
+        info = metadata.get(name)
+
+        if not info:
+            raise ValueError(
+                f"Tensor {name} has no metadata - is this a Lora safetensor?"
+            )
+
+        # Skip Textual Inversion embeds
+        if info == EMBED_FLAG:
+            continue
+
+        # Handle Loras
+        # Extract the targets
+        target = json.loads(info)
+
+        # Build the result lists - Python needs us to preallocate lists to insert into them
+        module_keys = list(module_keys)
+        ranks = [4] * (len(module_keys) // 2)
+        weights = [None] * len(module_keys)
+
+        for key in module_keys:
+            # Split the model name and index out of the key
+            _, idx, direction = key.split(":")
+            idx = int(idx)
+
+            # Add the rank
+            ranks[idx] = int(metadata[f"{name}:{idx}:rank"])
+
+            # Insert the weight into the list
+            idx = idx * 2 + (1 if direction == "down" else 0)
+            weights[idx] = nn.parameter.Parameter(safeloras.get_tensor(key))
+
+        loras[name] = (weights, ranks, target)
+
+    return loras
+
+
+def parse_safeloras_embeds(
+    safeloras,
+) -> Dict[str, torch.Tensor]:
+    """
+    Converts a loaded safetensor file that contains Textual Inversion embeds into
+    a dictionary of embed_token: Tensor
+    """
+    embeds = {}
+    metadata = safeloras.metadata()
+
+    for key in safeloras.keys():
+        # Only handle Textual Inversion embeds
+        meta = metadata.get(key)
+        if not meta or meta != EMBED_FLAG:
+            continue
+
+        embeds[key] = safeloras.get_tensor(key)
+
+    return embeds
+
+
+def load_safeloras(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras(safeloras)
+
+
+def load_safeloras_embeds(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras_embeds(safeloras)
+
+
+def load_safeloras_both(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras(safeloras), parse_safeloras_embeds(safeloras)
+
+
+def collapse_lora(model, alpha=1.0):
+
+    for _module, name, _child_module in _find_modules(
+        model,
+        UNET_EXTENDED_TARGET_REPLACE | TEXT_ENCODER_EXTENDED_TARGET_REPLACE,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
+    ):
+
+        if isinstance(_child_module, LoraInjectedLinear):
+            print("Collapsing Lin Lora in", name)
+
+            _child_module.linear.weight = nn.Parameter(
+                _child_module.linear.weight.data
+                + alpha
+                * (
+                    _child_module.lora_up.weight.data
+                    @ _child_module.lora_down.weight.data
+                )
+                .type(_child_module.linear.weight.dtype)
+                .to(_child_module.linear.weight.device)
+            )
+
+        else:
+            print("Collapsing Conv Lora in", name)
+            _child_module.conv.weight = nn.Parameter(
+                _child_module.conv.weight.data
+                + alpha
+                * (
+                    _child_module.lora_up.weight.data.flatten(start_dim=1)
+                    @ _child_module.lora_down.weight.data.flatten(start_dim=1)
+                )
+                .reshape(_child_module.conv.weight.data.shape)
+                .type(_child_module.conv.weight.dtype)
+                .to(_child_module.conv.weight.device)
+            )
+
+
+def monkeypatch_or_replace_lora(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    r: Union[int, List[int]] = 4,
+):
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[nn.Linear, LoraInjectedLinear]
+    ):
+        _source = (
+            _child_module.linear
+            if isinstance(_child_module, LoraInjectedLinear)
+            else _child_module
+        )
+
+        weight = _source.weight
+        bias = _source.bias
+        _tmp = LoraInjectedLinear(
+            _source.in_features,
+            _source.out_features,
+            _source.bias is not None,
+            r=r.pop(0) if isinstance(r, list) else r,
+        )
+        _tmp.linear.weight = weight
+
+        if bias is not None:
+            _tmp.linear.bias = bias
+
+        # switch the module
+        _module._modules[name] = _tmp
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype)
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype)
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def monkeypatch_or_replace_lora_extended(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    r: Union[int, List[int]] = 4,
+):
+    for _module, name, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[nn.Linear, LoraInjectedLinear, nn.Conv2d, LoraInjectedConv2d],
+    ):
+
+        if (_child_module.__class__ == nn.Linear) or (
+            _child_module.__class__ == LoraInjectedLinear
+        ):
+            if len(loras[0].shape) != 2:
+                continue
+
+            _source = (
+                _child_module.linear
+                if isinstance(_child_module, LoraInjectedLinear)
+                else _child_module
+            )
+
+            weight = _source.weight
+            bias = _source.bias
+            _tmp = LoraInjectedLinear(
+                _source.in_features,
+                _source.out_features,
+                _source.bias is not None,
+                r=r.pop(0) if isinstance(r, list) else r,
+            )
+            _tmp.linear.weight = weight
+
+            if bias is not None:
+                _tmp.linear.bias = bias
+
+        elif (_child_module.__class__ == nn.Conv2d) or (
+            _child_module.__class__ == LoraInjectedConv2d
+        ):
+            if len(loras[0].shape) != 4:
+                continue
+            _source = (
+                _child_module.conv
+                if isinstance(_child_module, LoraInjectedConv2d)
+                else _child_module
+            )
+
+            weight = _source.weight
+            bias = _source.bias
+            _tmp = LoraInjectedConv2d(
+                _source.in_channels,
+                _source.out_channels,
+                _source.kernel_size,
+                _source.stride,
+                _source.padding,
+                _source.dilation,
+                _source.groups,
+                _source.bias is not None,
+                r=r.pop(0) if isinstance(r, list) else r,
+            )
+
+            _tmp.conv.weight = weight
+
+            if bias is not None:
+                _tmp.conv.bias = bias
+
+        # switch the module
+        _module._modules[name] = _tmp
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype)
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype)
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def monkeypatch_or_replace_safeloras(models, safeloras):
+    loras = parse_safeloras(safeloras)
+
+    for name, (lora, ranks, target) in loras.items():
+        model = getattr(models, name, None)
+
+        if not model:
+            print(f"No model provided for {name}, contained in Lora")
+            continue
+
+        monkeypatch_or_replace_lora_extended(model, lora, target, ranks)
+
+
+def monkeypatch_remove_lora(model):
+    for _module, name, _child_module in _find_modules(
+        model, search_class=[LoraInjectedLinear, LoraInjectedConv2d]
+    ):
+        if isinstance(_child_module, LoraInjectedLinear):
+            _source = _child_module.linear
+            weight, bias = _source.weight, _source.bias
+
+            _tmp = nn.Linear(
+                _source.in_features, _source.out_features, bias is not None
+            )
+
+            _tmp.weight = weight
+            if bias is not None:
+                _tmp.bias = bias
+
+        else:
+            _source = _child_module.conv
+            weight, bias = _source.weight, _source.bias
+
+            _tmp = nn.Conv2d(
+                in_channels=_source.in_channels,
+                out_channels=_source.out_channels,
+                kernel_size=_source.kernel_size,
+                stride=_source.stride,
+                padding=_source.padding,
+                dilation=_source.dilation,
+                groups=_source.groups,
+                bias=bias is not None,
+            )
+
+            _tmp.weight = weight
+            if bias is not None:
+                _tmp.bias = bias
+
+        _module._modules[name] = _tmp
+
+
+def monkeypatch_add_lora(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    alpha: float = 1.0,
+    beta: float = 1.0,
+):
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[LoraInjectedLinear]
+    ):
+        weight = _child_module.linear.weight
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype).to(weight.device) * alpha
+            + _module._modules[name].lora_up.weight.to(weight.device) * beta
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype).to(weight.device) * alpha
+            + _module._modules[name].lora_down.weight.to(weight.device) * beta
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def tune_lora_scale(model, alpha: float = 1.0):
+    for _module in model.modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
+            _module.scale = alpha
+
+
+def set_lora_diag(model, diag: torch.Tensor):
+    for _module in model.modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
+            _module.set_selector_from_diag(diag)
+
+
+def _text_lora_path(path: str) -> str:
+    assert path.endswith(".pt"), "Only .pt files are supported"
+    return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"])
+
+
+def _ti_lora_path(path: str) -> str:
+    assert path.endswith(".pt"), "Only .pt files are supported"
+    return ".".join(path.split(".")[:-1] + ["ti", "pt"])
+
+
+def apply_learned_embed_in_clip(
+    learned_embeds,
+    text_encoder,
+    tokenizer,
+    token: Optional[Union[str, List[str]]] = None,
+    idempotent=False,
+):
+    if isinstance(token, str):
+        trained_tokens = [token]
+    elif isinstance(token, list):
+        assert len(learned_embeds.keys()) == len(
+            token
+        ), "The number of tokens and the number of embeds should be the same"
+        trained_tokens = token
+    else:
+        trained_tokens = list(learned_embeds.keys())
+
+    for token in trained_tokens:
+        print(token)
+        embeds = learned_embeds[token]
+
+        # cast to dtype of text_encoder
+        dtype = text_encoder.get_input_embeddings().weight.dtype
+        num_added_tokens = tokenizer.add_tokens(token)
+
+        i = 1
+        if not idempotent:
+            while num_added_tokens == 0:
+                print(f"The tokenizer already contains the token {token}.")
+                token = f"{token[:-1]}-{i}>"
+                print(f"Attempting to add the token {token}.")
+                num_added_tokens = tokenizer.add_tokens(token)
+                i += 1
+        elif num_added_tokens == 0 and idempotent:
+            print(f"The tokenizer already contains the token {token}.")
+            print(f"Replacing {token} embedding.")
+
+        # resize the token embeddings
+        text_encoder.resize_token_embeddings(len(tokenizer))
+
+        # get the id for the token and assign the embeds
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        text_encoder.get_input_embeddings().weight.data[token_id] = embeds
+    return token
+
+
+def load_learned_embed_in_clip(
+    learned_embeds_path,
+    text_encoder,
+    tokenizer,
+    token: Optional[Union[str, List[str]]] = None,
+    idempotent=False,
+):
+    learned_embeds = torch.load(learned_embeds_path)
+    apply_learned_embed_in_clip(
+        learned_embeds, text_encoder, tokenizer, token, idempotent
+    )
+
+
+def patch_pipe(
+    pipe,
+    maybe_unet_path,
+    token: Optional[str] = None,
+    r: int = 4,
+    patch_unet=True,
+    patch_text=True,
+    patch_ti=True,
+    idempotent_token=True,
+    unet_target_replace_module=DEFAULT_TARGET_REPLACE,
+    text_target_replace_module=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
+):
+    if maybe_unet_path.endswith(".pt"):
+        # torch format
+
+        if maybe_unet_path.endswith(".ti.pt"):
+            unet_path = maybe_unet_path[:-6] + ".pt"
+        elif maybe_unet_path.endswith(".text_encoder.pt"):
+            unet_path = maybe_unet_path[:-16] + ".pt"
+        else:
+            unet_path = maybe_unet_path
+
+        ti_path = _ti_lora_path(unet_path)
+        text_path = _text_lora_path(unet_path)
+
+        if patch_unet:
+            print("LoRA : Patching Unet")
+            monkeypatch_or_replace_lora(
+                pipe.unet,
+                torch.load(unet_path),
+                r=r,
+                target_replace_module=unet_target_replace_module,
+            )
+
+        if patch_text:
+            print("LoRA : Patching text encoder")
+            monkeypatch_or_replace_lora(
+                pipe.text_encoder,
+                torch.load(text_path),
+                target_replace_module=text_target_replace_module,
+                r=r,
+            )
+        if patch_ti:
+            print("LoRA : Patching token input")
+            token = load_learned_embed_in_clip(
+                ti_path,
+                pipe.text_encoder,
+                pipe.tokenizer,
+                token=token,
+                idempotent=idempotent_token,
+            )
+
+    elif maybe_unet_path.endswith(".safetensors"):
+        safeloras = safe_open(maybe_unet_path, framework="pt", device="cpu")
+        monkeypatch_or_replace_safeloras(pipe, safeloras)
+        tok_dict = parse_safeloras_embeds(safeloras)
+        if patch_ti:
+            apply_learned_embed_in_clip(
+                tok_dict,
+                pipe.text_encoder,
+                pipe.tokenizer,
+                token=token,
+                idempotent=idempotent_token,
+            )
+        return tok_dict
+
+
+@torch.no_grad()
+def inspect_lora(model):
+    moved = {}
+
+    for name, _module in model.named_modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
+            ups = _module.lora_up.weight.data.clone()
+            downs = _module.lora_down.weight.data.clone()
+
+            wght: torch.Tensor = ups.flatten(1) @ downs.flatten(1)
+
+            dist = wght.flatten().abs().mean().item()
+            if name in moved:
+                moved[name].append(dist)
+            else:
+                moved[name] = [dist]
+
+    return moved
+
+
+def save_all(
+    unet,
+    text_encoder,
+    save_path,
+    placeholder_token_ids=None,
+    placeholder_tokens=None,
+    save_lora=True,
+    save_ti=True,
+    target_replace_module_text=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
+    target_replace_module_unet=DEFAULT_TARGET_REPLACE,
+    safe_form=True,
+):
+    if not safe_form:
+        # save ti
+        if save_ti:
+            ti_path = _ti_lora_path(save_path)
+            learned_embeds_dict = {}
+            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
+                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
+                print(
+                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
+                    learned_embeds[:4],
+                )
+                learned_embeds_dict[tok] = learned_embeds.detach().cpu()
+
+            torch.save(learned_embeds_dict, ti_path)
+            print("Ti saved to ", ti_path)
+
+        # save text encoder
+        if save_lora:
+
+            save_lora_weight(
+                unet, save_path, target_replace_module=target_replace_module_unet
+            )
+            print("Unet saved to ", save_path)
+
+            save_lora_weight(
+                text_encoder,
+                _text_lora_path(save_path),
+                target_replace_module=target_replace_module_text,
+            )
+            print("Text Encoder saved to ", _text_lora_path(save_path))
+
+    else:
+        assert save_path.endswith(
+            ".safetensors"
+        ), f"Save path : {save_path} should end with .safetensors"
+
+        loras = {}
+        embeds = {}
+
+        if save_lora:
+
+            loras["unet"] = (unet, target_replace_module_unet)
+            loras["text_encoder"] = (text_encoder, target_replace_module_text)
+
+        if save_ti:
+            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
+                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
+                print(
+                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
+                    learned_embeds[:4],
+                )
+                embeds[tok] = learned_embeds.detach().cpu()
+
+        save_safeloras_with_embeds(loras, embeds, save_path)

README.md

@@ -1,3 +1,511 @@
+<div align="center">
+
+# TherA: Thermal-Aware Visual-Language Prompting for<br>Controllable RGB-to-Thermal Infrared Translation
+
+[**Dong-Guw Lee**](https://scholar.google.com/citations?user=u6VDnlgAAAAJ&hl=ko)<sup>1*</sup>&emsp;
+[**Tai Hyoung Rhee**](https://scholar.google.com/citations?user=PF8EfdYAAAAJ&hl=en&oi=ao)<sup>1*</sup>&emsp;
+[**Hyunsoo Jang**](https://rpm.snu.ac.kr)<sup>1</sup><br>
+[**Young-Sik Shin**](https://scholar.google.com/citations?user=gGfBRawAAAAJ&hl=en&oi=ao)<sup>2</sup>&emsp;
+[**Ukcheol Shin**](https://scholar.google.com/citations?user=ZvxI80EAAAAJ&hl=ko&oi=ao)<sup>3</sup>&emsp;
+[**Ayoung Kim**](https://ayoungk.github.io/)<sup>1&dagger;</sup>
+
+<sup>1</sup>Seoul National University&emsp;
+<sup>2</sup>Kyungpook National University&emsp;
+<sup>3</sup>KENTECH  
+<sup>*</sup> Equal Contribution&emsp;
+<sup>&dagger;</sup> Corresponding Author
+
+**CVPR 2026**
+
+[![Project Page](https://img.shields.io/badge/Project_Page-TherA-blue)](https://donkeymouse.github.io/thera_cvpr26/)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.19430-b31b1b.svg)](https://arxiv.org/abs/2602.19430)
+[![GitHub](https://img.shields.io/badge/GitHub-donkeymouse%2FTherA-black)](https://github.com/donkeymouse/TherA)
+[![Weights](https://img.shields.io/badge/HuggingFace-Weights-yellow)](https://huggingface.co/donkeymouse/TherA/tree/main)
+[![Dataset](https://img.shields.io/badge/HuggingFace-R2T2-orange)](https://huggingface.co/datasets/donkeymouse/TherA-R2T2)
+[![Docker](https://img.shields.io/badge/Docker-donkeymouse%2Fthera-2496ED)](https://hub.docker.com/r/donkeymouse/thera)
+
+</div>
+
+<p align="center">
+  <img src="assets/method.png" width="90%" alt="TherA method overview">
+</p>
+
+
 ---
-license: mit
+
+## News
+
+- **2026-04-03**: TherA github repo opening
+- **2026-05-22**: TherA inference code and R2T2 dataset release.
+
 ---
+
+## Overview
+
+**TherA** is a controllable RGB-to-thermal infrared translation framework. Given an RGB image, TherA synthesizes a long-wave thermal infrared image using a latent-diffusion translator conditioned on thermal-aware visual-language features.
+
+TherA is designed for:
+
+- **RGB → TIR translation** for thermal perception research.
+- **Thermal-aware VLM conditioning** using LLaVA hidden-state features.
+- **Scene- and object-level controllability** across weather, time of day, and object state.
+- **Reference-cache inference**, allowing deployment without loading LLaVA at runtime.
+
+
+<div align="center">
+  <a href="https://www.youtube.com/watch?v=X60UxjGKQkg">
+    <img src="https://img.youtube.com/vi/X60UxjGKQkg/0.jpg" alt="TherA demo video" width="720">
+  </a>
+</div>
+
+
+---
+
+## Key Idea
+
+TherA does **not** condition directly on raw text during diffusion inference. Instead, it uses a **4096-dimensional LLaVA hidden state**, either:
+
+1. loaded from a precomputed `.pt` reference cache, or
+2. extracted on the fly using LLaVA.
+
+
+
+For resource limited environments, we recommend **reference-cache mode**. This mode uses precomputed LLaVA features such as `SUNNY.pt`, `CLOUDY.pt`, `RAINY.pt`, or `NIGHT.pt`, and therefore does **not** require loading LLaVA weights at runtime. An alternative would be to compute pre-computed LLaVA feature first followed by inferencing with reference-cache mode (upcoming feature). 
+
+---
+## Repository Layout
+
+```text
+TherA/
+├── infer_custom.py             # Batch RGB → TIR inference on a folder
+├── infer_example_guided.py     # Single-image / example-guided inference
+├── infer_palette.sh        # Run multiple weather/style palettes
+├── lavi_ip2p/                  # UNet 8-channel + adapter wrapper
+├── LaVi-Bridge/modules/        # TextAdapter architecture
+├── llava/                      # LLaVA code, only needed for on-the-fly mode
+├── thera_paths.py              # Default local weight paths
+├── thera_llava.py              # Lazy LLaVA loader
+└── weights/                    # Download weights here; not tracked by git
+    ├── model.pt                # TherA Model
+    ├── merged_models/          # Initialization model
+    │   ├── unet/
+    │   └── adapter/
+    ├── stable-diffusion/       
+    │   ├── vae/
+    │   └── scheduler/
+    ├── reference_caches/
+    │   ├── SUNNY.pt
+    │   ├── CLOUDY.pt
+    │   ├── RAINY.pt
+    │   └── NIGHT.pt
+    ├── reference_caches/
+    │   │   ├── SUNNY.pt
+    │   │   ├── CLOUDY.pt
+    │   │   ├── RAINY.pt
+    │   │   └── NIGHT.pt
+    └── TherA-VLM/                  # Optional; only for on-the-fly mode
+        ├── adaptor_config.json/
+        └── adapter_model.safetensors
+        └── config.json
+        └── non_lora_trainables.bin
+        └── trainer_state.json
+```
+
+---
+
+## Installation
+
+### Option 1: Local Python Environment
+
+```bash
+git clone https://github.com/donkeymouse/TherA.git
+cd TherA
+
+python -m venv .venv
+source .venv/bin/activate
+
+pip install --upgrade pip
+pip install -r requirements.txt
+```
+
+**Recommended environment**
+
+- Python 3.10+
+- CUDA-capable GPU
+- 16 GB+ VRAM recommended for comfortable inference
+
+---
+
+### Option 2: Docker
+
+A prebuilt Docker image is available at:
+
+```bash
+docker pull donkeymouse/thera:latest
+```
+
+Example interactive run:
+
+```bash
+docker run --gpus all --rm -it \
+  -v "$(pwd)":/workspace/TherA \
+  -w /workspace/TherA \
+  donkeymouse/thera:latest \
+  bash
+```
+
+Then run inference commands from inside the container.
+
+---
+
+## Download Weights
+
+TherA weights are hosted on Hugging Face:
+
+```bash
+pip install -U huggingface_hub
+
+huggingface-cli download donkeymouse/TherA \
+  --local-dir weights
+```
+
+After downloading, your `weights/` directory should contain:
+
+| Path | Description | Required? |
+|---|---|---|
+| `weights/model.pt` | TherA trained UNet and adapter checkpoint | Yes |
+| `weights/merged_models/unet/` | UNet architecture/config files | Yes |
+| `weights/merged_models/adapter/` | TextAdapter architecture/config files | Yes |
+| `weights/stable-diffusion/vae/` | Stable Diffusion VAE | Yes |
+| `weights/stable-diffusion/scheduler/` | DDIM scheduler config | Yes |
+| `weights/reference_caches/*.pt` | Precomputed LLaVA hidden states for inference palettes | Recommended |
+| `weights/TherA-VLM/` | LLaVA weights for on-the-fly feature extraction | Optional |
+
+---
+
+## Optional: Download LLaVA Weights
+
+LLaVA is only required for **on-the-fly feature extraction** or **two-image guided mode**. It is not required for reference-cache inference.
+
+```bash
+huggingface-cli download llava-hf/llava-1.5-7b-hf \
+  --local-dir weights/llava-1.5-7b-hf
+```
+
+
+---
+
+## Quick Start
+
+---
+
+## Full RGB-TIR translation using TherA-VLM
+
+Use this mode if you want to extract hidden states from TherA directly at runtime from an RGB image and prompt.
+
+```bash
+python infer_custom.py \
+  --rgb-dir examples/rgb \
+  --output-dir preds \
+  --llava-base-path weights/llava-1.5-7b-hf \
+  --llava-lora-path weights/TherA-VLM \
+  --llava-prompt "How would this RGB scene appear in long-wave thermal infrared spectrum."
+```
+
+This mode is more expensive because it loads LLaVA during inference.
+
+
+
+### Reference-Guided Image Translation Mode
+
+This mode extracts LLaVA features from a reference RGB image and applies them to a target RGB image.
+
+```bash
+python infer_example_guided.py \
+  --mode two-image \
+  --reference-image examples/ref/rgb.jpg \
+  --input-image examples/rgb/scene.jpg \
+  --output preds/scene_tir.png \
+  --llava-base-path weights/llava-1.5-7b-hf \
+  --llava-lora-path weights/TherA-VLM
+```
+
+---
+
+### Recursive Folder Inference
+
+```bash
+python infer_custom.py \
+  --rgb-dir /path/to/dataset/RGB \
+  --output-dir preds \
+  --reference-cache weights/reference_caches/SUNNY.pt \
+  --recursive
+```
+
+When `--recursive` is used, the output folder preserves the input directory structure.
+
+---
+---
+
+## Reference-cache Mode
+Reference-cache mode is the recommended if you are lacking GPU memory. It does not load LLaVA at runtime.
+
+```bash
+python infer_custom.py \
+  --rgb-dir examples/rgb \
+  --output-dir preds/sunny \
+  --reference-cache weights/reference_caches/SUNNY.pt
+```
+
+The script reads all images in `examples/rgb` and writes translated TIR images to `preds/sunny`.
+
+A lighter version of the text-guided image translation module. 
+
+Example palette caches:
+
+```text
+weights/reference_caches/SUNNY.pt
+weights/reference_caches/CLOUDY.pt
+weights/reference_caches/RAINY.pt
+weights/reference_caches/NIGHT.pt
+```
+
+You can use different pallete cache to achieve different translation effects. 
+
+---
+
+## Inference Modes
+
+| Mode | Main flag / script | LLaVA weights needed? | Recommended use |
+|---|---|---:|---|
+| Reference cache | `--reference-cache path.pt` | No | Default deployment and fast inference |
+| Per-image cache directory | `--cache-dir dir/` | No | Precomputed feature per image |
+| Full RGB-TIR translation| `--llava-base-path ...` | Yes | Runtime prompt/image conditioning |
+| Reference image-guided translation| `infer_example_guided.py --mode two-image` | Yes | Apply reference-image conditioning |
+
+---
+
+## Reference Cache Format
+
+Reference caches are precomputed LLaVA hidden states saved as `.pt` files.
+
+Supported tensor shapes:
+
+```text
+[1, L, 4096]
+[L, 4096]
+```
+
+A single reference cache can be applied to all input images as a global thermal/weather/style condition.
+
+---
+
+## Important Arguments
+
+| Argument | Default | Description |
+|---|---:|---|
+| `--checkpoint` | `weights/checkpoint` | Directory containing `model.pt` |
+| `--merged-model-path` | `weights/merged_models` | Directory containing UNet and adapter configs |
+| `--pretrained-sd` | `weights/stable-diffusion` | Directory containing VAE and scheduler |
+| `--rgb-dir` | Required | Folder of RGB images for batch inference |
+| `--output-dir` | `custom_predictions` | Output folder for predictions |
+| `--reference-cache` | `None` | Single `.pt` cache used for all images |
+| `--cache-dir` | `None` | Folder of per-image `.pt` caches matched by filename stem |
+| `--llava-base-path` | `None` | Base LLaVA model path for on-the-fly mode |
+| `--llava-lora-path` | `None` | Optional LLaVA LoRA path |
+| `--llava-prompt` | thermal prompt | Prompt used for default inference/text-guided translation|
+| `--num-steps` | `100` | DDIM sampling steps |
+| `--cfg-text` | `3.5` | Text/VLM guidance strength |
+| `--cfg-image` | `1.5` | Image guidance strength |
+| `--target-size` | Auto | Resize image to this square size; otherwise dimensions are rounded to multiples of 32 |
+| `--recursive` | Off | Recursively process subdirectories |
+| `--device` | `cuda` | Device for inference |
+
+---
+
+## Architecture
+
+```text
+RGB image
+   │
+   ▼
+VAE encoder ──► RGB latents ───────────────────┐
+                                                │
+                                                ├──► 8-channel diffusion UNet ──► VAE decoder ──► TIR image
+                                                │
+LLaVA hidden state, 4096-d ──► TextAdapter ─────┘
+                              768-d cross-attention tokens
+```
+
+TherA uses dual classifier-free guidance at inference by combining:
+
+- full conditioning,
+- image-only conditioning,
+- text/VLM-only conditioning.
+
+---
+
+## R2T2 Dataset
+
+TherA is trained with **R2T2**, a large-scale RGB–TIR–Text dataset.
+
+R2T2 includes:
+
+- **112,970 aligned triplets**: RGB image, TIR image, and canonical thermal schema.
+- Scene diversity across driving, CCTV, aerial, and ego-view settings.
+- Temporal diversity across day/night and diurnal transitions.
+- Environmental diversity across weather, season, and illumination.
+- Material- and object-level annotations with structured canonicalization.
+- Data compiled from multiple aligned RGB–TIR datasets with additional pseudo-aligned pairs.
+
+Dataset page:
+
+```text
+https://huggingface.co/datasets/donkeymouse/TherA-R2T2
+```
+
+Example structure:
+
+```text
+R2T2/
+├── ${DATASET_NAME}/
+│   └── ${SEQUENCE_NAME}/
+│       ├── RGB/
+│       │   ├── 1.jpg
+│       │   └── ...
+│       └── TIR/
+│           ├── 1.jpg
+│           └── ...
+├── ViVID/
+│   ├── img_campus_day1/
+│   │   ├── RGB/
+│   │   │   ├── 000001.png
+│   │   │   └── ...
+│   │   └── TIR/
+│   │       ├── 000001.png
+│   │       └── ...
+│   └── ...
+└── ...
+```
+
+---
+
+## Troubleshooting
+
+### `Checkpoint not found: weights/model.pt`
+
+Download the TherA weights and make sure `model.pt` is located at:
+
+```text
+weights/model.pt
+```
+
+---
+
+### `OSError: ... stable-diffusion/vae`
+
+Make sure the Stable Diffusion VAE and scheduler folders are present:
+
+```text
+weights/stable-diffusion/vae/
+weights/stable-diffusion/scheduler/
+```
+
+---
+
+### Outputs look identical across palettes
+
+Try increasing text/VLM guidance:
+
+```bash
+--cfg-text 7.5
+```
+
+Also verify that your reference cache files are distinct:
+
+```text
+SUNNY.pt
+CLOUDY.pt
+RAINY.pt
+NIGHT.pt
+```
+
+---
+
+### CUDA out of memory
+
+Try reducing the image size:
+
+```bash
+--target-size 512
+```
+
+You can also run one image at a time with:
+
+```bash
+python infer_example_guided.py
+```
+
+---
+
+### LLaVA import or loading errors
+
+Use reference-cache mode if you do not need runtime LLaVA extraction:
+
+```bash
+--reference-cache weights/reference_caches/SUNNY.pt
+```
+
+For on-the-fly mode, make sure the LLaVA base model and TherA weights are correctly loaded`.
+
+---
+
+## TODOs
+- [x] inference code and R2T2 dataset
+- [] Upload cache extraction code
+- [] Improve text-guidance 
+
+
+## Citation
+
+If you find TherA useful for your research, please cite:
+
+```bibtex
+@inproceedings{lee2026thera,
+  title     = {TherA: Thermal-Aware Visual-Language Prompting for Controllable RGB-to-Thermal Infrared Translation},
+  author    = {Lee, Dong-Guw and Rhee, Tai Hyoung and Jang, Hyunsoo and Shin, Young-Sik and Shin, Ukcheol and Kim, Ayoung},
+  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  year      = {2026}
+}
+```
+
+You may also cite the arXiv version:
+
+```bibtex
+@article{lee2026thera_arxiv,
+  title   = {TherA: Thermal-Aware Visual-Language Prompting for Controllable RGB-to-Thermal Infrared Translation},
+  author  = {Lee, Dong-Guw and Rhee, Tai Hyoung and Jang, Hyunsoo and Shin, Young-Sik and Shin, Ukcheol and Kim, Ayoung},
+  journal = {arXiv preprint arXiv:2602.19430},
+  year    = {2026}
+}
+```
+
+---
+
+## Acknowledgements
+
+TherA builds on open-source components from the vision-language and diffusion communities, including LLaVA, Stable Diffusion, Diffusers, and LaVi-Bridge-style adapter architectures.
+
+---
+
+## License
+
+See `LICENSE` for details.
+
+Third-party models, datasets, and libraries retain their own licenses. Please review the licenses for LLaVA, Stable Diffusion, Hugging Face model files, and any external datasets before use.
+
+## Contact
+If you have any questions, contact here please
+```
+donkeymouse@snu.ac.kr
+```

infer_custom.py

@@ -0,0 +1,526 @@
+#!/usr/bin/env python3
+"""
+Batch RGB → TIR inference for TherA.
+
+Recommended (no LLaVA weights needed): use a reference .pt cache
+  python infer_custom.py --rgb-dir ./examples/rgb --output-dir ./preds \\
+      --reference-cache weights/reference_caches/SUNNY.pt
+
+On-the-fly LLaVA (optional): requires separate LLaVA checkpoints
+  python infer_custom.py --rgb-dir ./examples/rgb --output-dir ./preds \\
+      --llava-base-path weights/llava/llava-1.5-7b-hf \\
+      --llava-lora-path weights/llava/llava-miragehd-lora
+
+Single-image example-guided mode: see infer_example_guided.py
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import List
+
+import torch
+from PIL import Image
+from tqdm import tqdm
+from contextlib import nullcontext
+
+from thera_paths import (
+    DEFAULT_CHECKPOINT,
+    DEFAULT_MERGED_MODEL,
+    DEFAULT_PRETRAINED_SD,
+    setup_project_path,
+)
+from thera_llava import create_frozen_llava_extractor
+
+setup_project_path()
+
+from diffusers import AutoencoderKL, DDIMScheduler, UNet2DConditionModel
+from torchvision import transforms
+from torchvision.utils import save_image
+
+from lavi_ip2p.llava_adapter import create_llava_adapter
+from lavi_ip2p.unet_8ch import convert_unet_to_8ch
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate TIR predictions from custom RGB images")
+    
+    # Model paths (defaults under weights/ — see README.md)
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default=str(DEFAULT_CHECKPOINT),
+        help="Directory containing model.pt (TherA UNet + adapter weights)",
+    )
+    parser.add_argument(
+        "--merged-model-path",
+        type=str,
+        default=str(DEFAULT_MERGED_MODEL),
+        help="Directory with unet/ and adapter/ architecture configs",
+    )
+    parser.add_argument(
+        "--pretrained-sd",
+        type=str,
+        default=str(DEFAULT_PRETRAINED_SD),
+        help="Stable Diffusion folder with vae/ and scheduler/ subfolders",
+    )
+    
+    # LLaVA paths (only for on-the-fly feature extraction)
+    parser.add_argument(
+        "--llava-base-path",
+        type=str,
+        default=None,
+        help="Base LLaVA model. Required unless --reference-cache or --cache-dir is set.",
+    )
+    parser.add_argument(
+        "--llava-lora-path",
+        type=str,
+        default=None,
+        help="Optional LLaVA LoRA weights (on-the-fly mode only)",
+    )
+    parser.add_argument("--llava-prompt", type=str, 
+                       default="How would this RGB scene appear in long-wave thermal infrared spectrum",
+                       help="Prompt for LLaVA feature extraction")
+    
+    # Reference cache mode (alternative to on-the-fly extraction)
+    parser.add_argument("--reference-cache", type=str, default=None,
+                       help="Path to a reference .pt cache file. If provided, uses this fixed feature for all images instead of extracting per-image.")
+    parser.add_argument("--cache-dir", type=str, default=None,
+                       help="Directory with per-image .pt cache files (matched by filename stem). Alternative to --reference-cache.")
+    
+    # Data
+    parser.add_argument("--rgb-dir", type=str, required=True,
+                       help="Directory containing RGB images")
+    parser.add_argument("--output-dir", type=str, default="custom_predictions",
+                       help="Output directory for predictions")
+    parser.add_argument("--recursive", action="store_true",
+                       help="Search for images recursively in subdirectories")
+    
+    # Sampling
+    parser.add_argument("--num-steps", type=int, default=100, help="DDIM sampling steps")
+    parser.add_argument("--cfg-text", type=float, default=3.5, help="CFG scale for text")
+    parser.add_argument("--cfg-image", type=float, default=1.5, help="CFG scale for image")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    
+    # Processing
+    parser.add_argument("--target-size", type=int, default=None,
+                       help="Resize images to this size (default: keep original, round to 32)")
+    parser.add_argument("--batch-size", type=int, default=1,
+                       help="Batch size (>1 for faster processing if you have enough VRAM)")
+    
+    # Device
+    parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu)")
+    
+    return parser.parse_args()
+
+
+def load_models(args):
+    """Load VAE, UNet, Adapter, Scheduler, and optionally LLaVA Extractor"""
+    print("\n" + "="*80)
+    print(f"Loading models on device {args.device}...")
+    print("="*80)
+    
+    device = torch.device(args.device)
+    
+    # Determine if we need LLaVA extractor
+    use_cache = args.reference_cache is not None or args.cache_dir is not None
+    
+    # 1. VAE
+    print("\n[1/5] Loading VAE...")
+    vae = AutoencoderKL.from_pretrained(args.pretrained_sd, subfolder="vae")
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(device)
+    print("✓ VAE loaded")
+    
+    # 2. Scheduler
+    print("\n[2/5] Loading scheduler...")
+    scheduler = DDIMScheduler.from_pretrained(args.pretrained_sd, subfolder="scheduler")
+    print("✓ Scheduler loaded")
+    
+    # 3. UNet (load from checkpoint)
+    checkpoint_dir = Path(args.checkpoint)
+    print(f"\n[3/5] Loading UNet from checkpoint: {checkpoint_dir}")
+    checkpoint_path = checkpoint_dir / "model.pt"
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+    
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    
+    # Load base UNet architecture from merged model
+    merged_unet_path = Path(args.merged_model_path) / "unet"
+    unet = UNet2DConditionModel.from_pretrained(str(merged_unet_path))
+    
+    # Convert to 8-channel for image-to-image task
+    print("  Converting UNet to 8-channel input...")
+    unet = convert_unet_to_8ch(unet)
+    
+    # Load trained weights
+    unet.load_state_dict(checkpoint['unet'])
+    unet.requires_grad_(False)
+    unet.eval()
+    unet.to(device)
+    
+    global_step = checkpoint.get('global_step', 'unknown')
+    epoch = checkpoint.get('epoch', 'unknown')
+    print(f"✓ UNet loaded (step={global_step}, epoch={epoch})")
+    
+    # 4. Adapter
+    print("\n[4/5] Loading LLaVA Adapter...")
+    merged_adapter_path = Path(args.merged_model_path) / "adapter"
+    adapter = create_llava_adapter(
+        adapter_path=str(merged_adapter_path),
+        use_rms_norm=False,
+        learnable_scale=False,
+        init_scale=1.0,
+        freeze_adapter=True,
+    )
+    
+    # Load trained adapter weights if available in checkpoint
+    if 'llava_adapter' in checkpoint:
+        adapter.load_state_dict(checkpoint['llava_adapter'])
+        print("✓ Adapter loaded from checkpoint")
+    else:
+        print("✓ Adapter loaded (using base weights)")
+    
+    adapter.requires_grad_(False)
+    adapter.eval()
+    adapter.to(device)
+    
+    # 5. LLaVA Extractor (optional - only if not using cache)
+    llava_extractor = None
+    if use_cache:
+        print("\n[5/5] Skipping LLaVA extractor (using cached features)")
+    else:
+        print("\n[5/5] Loading LLaVA Feature Extractor...")
+        if not args.llava_base_path:
+            raise ValueError("--llava-base-path is required when not using --reference-cache or --cache-dir")
+        llava_extractor = create_frozen_llava_extractor(
+            llava_base_path=args.llava_base_path,
+            llava_lora_path=args.llava_lora_path,
+            device=args.device,
+            load_8bit=False,
+            load_4bit=False,
+            merge_lora=True,
+        )
+        print("✓ LLaVA extractor loaded")
+    
+    return vae, unet, adapter, scheduler, llava_extractor, device
+
+
+def find_images(rgb_dir: Path, recursive: bool = False) -> List[Path]:
+    """Find all image files in directory"""
+    if recursive:
+        images = []
+        for ext in IMAGE_EXTENSIONS:
+            images.extend(rgb_dir.rglob(f"*{ext}"))
+            images.extend(rgb_dir.rglob(f"*{ext.upper()}"))
+        return sorted(set(images))
+    else:
+        images = []
+        for ext in IMAGE_EXTENSIONS:
+            images.extend(rgb_dir.glob(f"*{ext}"))
+            images.extend(rgb_dir.glob(f"*{ext.upper()}"))
+        return sorted(set(images))
+
+
+def load_and_prepare_image(image_path: Path, target_size: int = None, device: torch.device = None):
+    """Load RGB image and prepare for processing"""
+    rgb_img = Image.open(image_path).convert('RGB')
+    rgb_w, rgb_h = rgb_img.size
+    
+    # Determine target dimensions
+    if target_size is not None:
+        target_w = target_h = target_size
+    else:
+        # Ensure dimensions divisible by 32 (required for VAE)
+        target_w = ((rgb_w + 31) // 32) * 32
+        target_h = ((rgb_h + 31) // 32) * 32
+    
+    # Resize if needed
+    if target_w != rgb_w or target_h != rgb_h:
+        rgb_img = rgb_img.resize((target_w, target_h), Image.Resampling.LANCZOS)
+    
+    # Convert to tensor for diffusion model
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),  # [-1, 1]
+    ])
+    
+    rgb_tensor = transform(rgb_img).unsqueeze(0)
+    if device is not None:
+        rgb_tensor = rgb_tensor.to(device)
+    
+    return rgb_tensor, rgb_img  # Return both tensor and PIL image
+
+
+def load_llava_cache(cache_path: Path, device: torch.device):
+    """Load a cached LLaVA hidden state from .pt file"""
+    try:
+        llava_hidden = torch.load(cache_path, map_location="cpu", weights_only=True)
+    except TypeError:
+        # For older torch versions without weights_only
+        llava_hidden = torch.load(cache_path, map_location="cpu")
+    
+    if not isinstance(llava_hidden, torch.Tensor):
+        llava_hidden = torch.tensor(llava_hidden)
+    
+    # Add batch dimension if needed
+    if llava_hidden.ndim == 2:
+        llava_hidden = llava_hidden.unsqueeze(0)
+    
+    return llava_hidden.to(device)
+
+
+def find_cache_for_image(image_path: Path, cache_dir: Path) -> Path:
+    """Find matching cache file for an image based on filename stem"""
+    stem = image_path.stem
+    for ext in ['.pt', '.pth']:
+        cache_path = cache_dir / f"{stem}{ext}"
+        if cache_path.exists():
+            return cache_path
+    raise FileNotFoundError(f"No cache file found for {image_path.name} in {cache_dir}")
+
+
+def create_null_text(llava_tokens):
+    """Create null tokens for unconditional generation"""
+    return torch.zeros_like(llava_tokens)
+
+
+@torch.no_grad()
+def translate_image(
+    vae,
+    unet,
+    adapter,
+    scheduler,
+    rgb_tensor,
+    llava_hidden,
+    num_steps=100,
+    cfg_text=3.5,
+    cfg_image=1.5,
+    device=torch.device("cuda"),
+):
+    """Translate RGB to TIR using LLaVA features + conditional generation with CFG"""
+    
+    # Encode RGB
+    rgb_latents = vae.encode(rgb_tensor).latent_dist.mode() * vae.config.scaling_factor
+    
+    # Process LLaVA features through adapter
+    use_autocast = device.type == "cuda"
+    llava_input = llava_hidden.to(device)
+    if use_autocast:
+        llava_input = llava_input.to(dtype=torch.bfloat16)
+    
+    autocast_ctx = (
+        torch.autocast(device_type=device.type, dtype=torch.bfloat16) 
+        if use_autocast else nullcontext()
+    )
+    with autocast_ctx:
+        llava_tokens = adapter(llava_input)
+    llava_tokens = llava_tokens.to(device=device, dtype=next(unet.parameters()).dtype)
+    
+    # Create null tokens for CFG
+    null_tokens = create_null_text(llava_tokens)
+    
+    # Initialize random noise
+    latent_shape = (1, 4, rgb_latents.shape[2], rgb_latents.shape[3])
+    latents = torch.randn(latent_shape, device=device, dtype=rgb_latents.dtype)
+    
+    # Setup scheduler
+    scheduler.set_timesteps(num_steps, device=device)
+    
+    zeros_rgb = torch.zeros_like(rgb_latents)
+    
+    for t in scheduler.timesteps:
+        t_batch = torch.full((1,), int(t), device=device, dtype=torch.long)
+        
+        # 3 forward passes for dual-CFG
+        x8_full = torch.cat([latents, rgb_latents], dim=1)
+        x8_no_img = torch.cat([latents, zeros_rgb], dim=1)
+        
+        # Full conditioning
+        eps_full = unet(x8_full, t_batch, encoder_hidden_states=llava_tokens).sample
+        
+        # No text (image only)
+        eps_no_text = unet(x8_full, t_batch, encoder_hidden_states=null_tokens).sample
+        
+        # No image (text only)
+        eps_no_img = unet(x8_no_img, t_batch, encoder_hidden_states=llava_tokens).sample
+        
+        # Combine with dual-CFG
+        eps = eps_full + cfg_text * (eps_full - eps_no_text) + cfg_image * (eps_full - eps_no_img)
+        eps = torch.clamp(eps, -5.0, 5.0)  # Safety guard
+        
+        latents = scheduler.step(eps, t, latents, eta=0.0).prev_sample
+    
+    # Decode
+    pred_tir = vae.decode(latents / vae.config.scaling_factor).sample
+    pred_tir = (pred_tir / 2 + 0.5).clamp(0, 1)
+    
+    return pred_tir
+
+
+def main():
+    args = parse_args()
+    
+    # Set seed
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+    
+    # Determine processing mode
+    use_reference = args.reference_cache is not None
+    use_cache_dir = args.cache_dir is not None
+    
+    if use_reference and use_cache_dir:
+        raise ValueError("Cannot use both --reference-cache and --cache-dir. Choose one.")
+    
+    print("\n" + "="*80)
+    if use_reference:
+        print("Custom Dataset Inference (Reference Cache Mode)")
+    elif use_cache_dir:
+        print("Custom Dataset Inference (Per-Image Cache Mode)")
+    else:
+        print("Custom Dataset Inference (On-the-fly LLaVA Features)")
+    print("="*80)
+    print(f"RGB directory: {args.rgb_dir}")
+    print(f"Checkpoint: {args.checkpoint}")
+    
+    if use_reference:
+        print(f"Reference cache: {args.reference_cache}")
+    elif use_cache_dir:
+        print(f"Cache directory: {args.cache_dir}")
+    else:
+        print(f"LLaVA base: {args.llava_base_path}")
+        if args.llava_lora_path:
+            print(f"LLaVA LoRA: {args.llava_lora_path}")
+        print(f"Prompt: {args.llava_prompt}")
+    
+    print(f"Output directory: {args.output_dir}")
+    print(f"DDIM steps: {args.num_steps}")
+    print(f"CFG scales: text={args.cfg_text}, image={args.cfg_image}")
+    print("="*80)
+    
+    # Create output directory
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Load models
+    vae, unet, adapter, scheduler, llava_extractor, device = load_models(args)
+    
+    # Load reference cache if provided
+    reference_hidden = None
+    if use_reference:
+        print(f"\nLoading reference cache from: {args.reference_cache}")
+        reference_hidden = load_llava_cache(Path(args.reference_cache), device)
+        print(f"✓ Reference cache loaded: {reference_hidden.shape}")
+    
+    # Validate cache directory if provided
+    cache_dir = None
+    if use_cache_dir:
+        cache_dir = Path(args.cache_dir)
+        if not cache_dir.exists():
+            raise FileNotFoundError(f"Cache directory not found: {cache_dir}")
+        print(f"\nUsing cache directory: {cache_dir}")
+    
+    # Find all RGB images
+    rgb_dir = Path(args.rgb_dir)
+    if not rgb_dir.exists():
+        raise FileNotFoundError(f"RGB directory not found: {rgb_dir}")
+    
+    image_files = find_images(rgb_dir, recursive=args.recursive)
+    print(f"\nFound {len(image_files)} images")
+    
+    if len(image_files) == 0:
+        print("No images found. Exiting.")
+        return
+    
+    # Process each image
+    print("\n" + "="*80)
+    print("Generating predictions...")
+    print("="*80)
+    
+    successful = 0
+    failed = 0
+    skipped = 0
+    
+    for img_path in tqdm(image_files, desc="Processing images"):
+        try:
+            # Load and prepare RGB image
+            rgb_tensor, rgb_pil = load_and_prepare_image(
+                img_path, 
+                target_size=args.target_size,
+                device=device
+            )
+            
+            # Get LLaVA hidden states based on mode
+            if use_reference:
+                # Use the same reference cache for all images
+                llava_hidden = reference_hidden
+            elif use_cache_dir:
+                # Find matching cache file for this image
+                try:
+                    cache_path = find_cache_for_image(img_path, cache_dir)
+                    llava_hidden = load_llava_cache(cache_path, device)
+                except FileNotFoundError as e:
+                    print(f"\nWarning: {e}")
+                    skipped += 1
+                    continue
+            else:
+                # Extract features on-the-fly using LLaVA
+                llava_hidden = llava_extractor.extract_hidden_states([rgb_pil], [args.llava_prompt])
+            
+            # Translate RGB to TIR
+            pred_tir = translate_image(
+                vae,
+                unet,
+                adapter,
+                scheduler,
+                rgb_tensor,
+                llava_hidden,
+                num_steps=args.num_steps,
+                cfg_text=args.cfg_text,
+                cfg_image=args.cfg_image,
+                device=device,
+            )
+            
+            # Save prediction
+            # Preserve relative path structure if recursive
+            if args.recursive:
+                rel_path = img_path.relative_to(rgb_dir)
+                save_path = output_dir / rel_path
+            else:
+                save_path = output_dir / img_path.name
+            
+            save_path.parent.mkdir(parents=True, exist_ok=True)
+            save_image(pred_tir.cpu(), str(save_path))
+            
+            successful += 1
+            
+        except Exception as e:
+            print(f"\nError processing {img_path}: {e}")
+            import traceback
+            traceback.print_exc()
+            failed += 1
+            continue
+    
+    # Print final results
+    print("\n" + "="*80)
+    print("INFERENCE COMPLETE")
+    print("="*80)
+    print(f"Total images: {len(image_files)}")
+    print(f"Successful: {successful}")
+    if skipped > 0:
+        print(f"Skipped (no cache): {skipped}")
+    print(f"Failed: {failed}")
+    if len(image_files) > 0:
+        print(f"Success rate: {successful/len(image_files)*100:.1f}%")
+    print(f"Predictions saved to: {args.output_dir}")
+    print("="*80)
+
+
+if __name__ == "__main__":
+    main()

infer_example_guided.py

@@ -0,0 +1,432 @@
+#!/usr/bin/env python3
+"""
+Single-image example-guided RGB → TIR translation.
+
+Mode: cached (recommended — no LLaVA weights needed)
+  python infer_example_guided.py --mode cached \\
+      --reference-cache weights/reference_caches/SUNNY.pt \\
+      --input-image examples/rgb/scene.jpg \\
+      --output preds/scene_tir.png
+
+Mode: two-image (extract features from a reference RGB on-the-fly)
+  python infer_example_guided.py --mode two-image \\
+      --reference-image examples/ref/rgb.jpg \\
+      --input-image examples/rgb/scene.jpg \\
+      --output preds/scene_tir.png \\
+      --llava-base-path weights/llava/llava-1.5-7b-hf \\
+      --llava-lora-path weights/llava/llava-miragehd-lora
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+from contextlib import nullcontext
+
+import torch
+from PIL import Image
+
+from thera_paths import (
+    DEFAULT_CHECKPOINT,
+    DEFAULT_MERGED_MODEL,
+    DEFAULT_PRETRAINED_SD,
+    setup_project_path,
+)
+from thera_llava import create_frozen_llava_extractor
+
+setup_project_path()
+
+from diffusers import AutoencoderKL, DDIMScheduler, UNet2DConditionModel
+from torchvision import transforms
+from torchvision.utils import save_image
+
+from lavi_ip2p.llava_adapter import create_llava_adapter
+from lavi_ip2p.unet_8ch import convert_unet_to_8ch
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Example-guided translation with two modes",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__
+    )
+    
+    # Mode selection
+    parser.add_argument("--mode", type=str, required=True, choices=["two-image", "cached"],
+                       help="Translation mode: 'two-image' (extract from reference image) or 'cached' (use .pt file)")
+    
+    # Input/Output
+    parser.add_argument("--input-image", type=str, required=True,
+                       help="Path to input RGB image (Image B) to be translated")
+    parser.add_argument("--output", type=str, required=True,
+                       help="Path to save output TIR image")
+    
+    # Reference (mode-dependent)
+    parser.add_argument("--reference-image", type=str, default=None,
+                       help="Path to reference RGB image (Image A). Required for 'two-image' mode.")
+    parser.add_argument("--reference-cache", type=str, default=None,
+                       help="Path to reference .pt cache file. Required for 'cached' mode.")
+    
+    # Model paths (defaults under weights/ — see README.md)
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default=str(DEFAULT_CHECKPOINT),
+        help="Directory containing model.pt",
+    )
+    parser.add_argument(
+        "--merged-model-path",
+        type=str,
+        default=str(DEFAULT_MERGED_MODEL),
+        help="Directory with unet/ and adapter/ architecture configs",
+    )
+    parser.add_argument(
+        "--pretrained-sd",
+        type=str,
+        default=str(DEFAULT_PRETRAINED_SD),
+        help="Stable Diffusion folder with vae/ and scheduler/ subfolders",
+    )
+    
+    # LLaVA paths (required for two-image mode)
+    parser.add_argument("--llava-base-path", type=str, default=None,
+                       help="Path to base LLaVA model. Required for 'two-image' mode.")
+    parser.add_argument("--llava-lora-path", type=str, default=None,
+                       help="Optional: Path to LLaVA LoRA weights")
+    parser.add_argument("--llava-prompt", type=str,
+                       default="How would this RGB scene appear in long-wave thermal infrared spectrum",
+                       help="Prompt for LLaVA feature extraction")
+    
+    # Sampling parameters
+    parser.add_argument("--num-steps", type=int, default=100, help="DDIM sampling steps")
+    parser.add_argument("--cfg-text", type=float, default=3.5, help="CFG scale for text")
+    parser.add_argument("--cfg-image", type=float, default=1.5, help="CFG scale for image")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    
+    # Processing
+    parser.add_argument("--target-size", type=int, default=None,
+                       help="Resize images to this size (default: keep original, round to 32)")
+    parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu)")
+    
+    return parser.parse_args()
+
+
+def validate_args(args):
+    """Validate argument combinations based on mode"""
+    if args.mode == "two-image":
+        if args.reference_image is None:
+            print("ERROR: --reference-image required for 'two-image' mode")
+            sys.exit(1)
+        if args.llava_base_path is None:
+            print("ERROR: --llava-base-path required for 'two-image' mode")
+            sys.exit(1)
+        if not Path(args.reference_image).exists():
+            print(f"ERROR: Reference image not found: {args.reference_image}")
+            sys.exit(1)
+    
+    elif args.mode == "cached":
+        if args.reference_cache is None:
+            print("ERROR: --reference-cache required for 'cached' mode")
+            sys.exit(1)
+        if not Path(args.reference_cache).exists():
+            print(f"ERROR: Reference cache not found: {args.reference_cache}")
+            sys.exit(1)
+    
+    if not Path(args.input_image).exists():
+        print(f"ERROR: Input image not found: {args.input_image}")
+        sys.exit(1)
+    
+    if not Path(args.checkpoint, "model.pt").exists():
+        print(f"ERROR: Checkpoint not found: {Path(args.checkpoint) / 'model.pt'}")
+        sys.exit(1)
+
+
+def load_models(args, need_llava=False):
+    """Load diffusion models and optionally LLaVA extractor"""
+    print("\n" + "="*80)
+    print(f"Loading models on device {args.device}...")
+    print("="*80)
+    
+    device = torch.device(args.device)
+    checkpoint_dir = Path(args.checkpoint)
+    
+    # 1. Load VAE
+    print("\n[1/5] Loading VAE...")
+    vae = AutoencoderKL.from_pretrained(args.pretrained_sd, subfolder="vae")
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(device)
+    print("✓ VAE loaded")
+    
+    # 2. Load Scheduler
+    print("\n[2/5] Loading scheduler...")
+    scheduler = DDIMScheduler.from_pretrained(args.pretrained_sd, subfolder="scheduler")
+    print("✓ Scheduler loaded")
+    
+    # 3. Load UNet (load from checkpoint)
+    print(f"\n[3/5] Loading UNet from checkpoint: {checkpoint_dir}")
+    checkpoint_path = checkpoint_dir / "model.pt"
+    if not checkpoint_path.exists():
+        print(f"ERROR: Checkpoint not found: {checkpoint_path}")
+        sys.exit(1)
+    
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    
+    # Load base UNet architecture from merged model
+    merged_unet_path = Path(args.merged_model_path) / "unet"
+    unet = UNet2DConditionModel.from_pretrained(str(merged_unet_path))
+    
+    # Convert to 8-channel for image-to-image task
+    print("  Converting UNet to 8-channel input...")
+    unet = convert_unet_to_8ch(unet)
+    
+    # Load trained weights
+    unet.load_state_dict(checkpoint['unet'])
+    unet.requires_grad_(False)
+    unet.eval()
+    unet.to(device)
+    
+    global_step = checkpoint.get('global_step', 'unknown')
+    epoch = checkpoint.get('epoch', 'unknown')
+    print(f"✓ UNet loaded (step={global_step}, epoch={epoch})")
+    
+    # 4. Load Adapter
+    print("\n[4/5] Loading LLaVA Adapter...")
+    merged_adapter_path = Path(args.merged_model_path) / "adapter"
+    adapter = create_llava_adapter(
+        adapter_path=str(merged_adapter_path),
+        use_rms_norm=False,
+        learnable_scale=False,
+        init_scale=1.0,
+        freeze_adapter=True,
+    )
+    
+    # Load trained adapter weights if available in checkpoint
+    if 'llava_adapter' in checkpoint:
+        adapter.load_state_dict(checkpoint['llava_adapter'])
+        print("✓ Adapter loaded from checkpoint")
+    else:
+        print("✓ Adapter loaded (using base weights)")
+    
+    adapter.requires_grad_(False)
+    adapter.eval()
+    adapter.to(device)
+    
+    # 5. Optionally load LLaVA
+    llava_extractor = None
+    if need_llava:
+        print("\n[5/5] Loading LLaVA Feature Extractor...")
+        if not args.llava_base_path:
+            print("ERROR: --llava-base-path is required for two-image mode")
+            sys.exit(1)
+        llava_extractor = create_frozen_llava_extractor(
+            llava_base_path=args.llava_base_path,
+            llava_lora_path=args.llava_lora_path,
+            device=args.device,
+            load_8bit=False,
+            load_4bit=False,
+            merge_lora=True,
+        )
+        print("✓ LLaVA extractor loaded")
+    else:
+        print("\n[5/5] Skipping LLaVA extractor (using cached features)")
+    
+    return vae, unet, adapter, scheduler, llava_extractor, device
+
+
+def load_and_prepare_image(image_path, target_size=None, device="cuda"):
+    """Load image and prepare tensor"""
+    rgb_img = Image.open(image_path).convert('RGB')
+    rgb_w, rgb_h = rgb_img.size
+    
+    # Determine target dimensions
+    if target_size is not None:
+        target_w = target_h = target_size
+    else:
+        # Ensure dimensions divisible by 32 (required for VAE)
+        target_w = ((rgb_w + 31) // 32) * 32
+        target_h = ((rgb_h + 31) // 32) * 32
+    
+    # Resize if needed
+    if target_w != rgb_w or target_h != rgb_h:
+        rgb_img = rgb_img.resize((target_w, target_h), Image.Resampling.LANCZOS)
+    
+    # Convert to tensor for diffusion model
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),  # [-1, 1]
+    ])
+    
+    rgb_tensor = transform(rgb_img).unsqueeze(0)
+    if device is not None:
+        rgb_tensor = rgb_tensor.to(device)
+    
+    return rgb_tensor, rgb_img  # Return both tensor and PIL image
+
+
+def load_llava_cache(cache_path, device):
+    """Load pre-cached LLaVA hidden states"""
+    try:
+        llava_hidden = torch.load(cache_path, map_location="cpu", weights_only=True)
+    except TypeError:
+        # For older torch versions without weights_only
+        llava_hidden = torch.load(cache_path, map_location="cpu")
+    
+    if not isinstance(llava_hidden, torch.Tensor):
+        llava_hidden = torch.tensor(llava_hidden)
+    
+    # Add batch dimension if needed
+    if llava_hidden.ndim == 2:
+        llava_hidden = llava_hidden.unsqueeze(0)
+    
+    return llava_hidden.to(device)
+
+
+@torch.no_grad()
+def translate_image(vae, unet, adapter, scheduler, rgb_tensor, llava_hidden, 
+                   num_steps=100, cfg_text=3.5, cfg_image=1.5, device="cuda"):
+    """Translate RGB to TIR using LLaVA features + conditional generation with CFG"""
+    
+    # Encode RGB
+    rgb_latents = vae.encode(rgb_tensor).latent_dist.mode() * vae.config.scaling_factor
+    batch_size = rgb_latents.shape[0]
+    
+    # Process LLaVA features through adapter
+    use_autocast = device.type == "cuda"
+    llava_input = llava_hidden.to(device)
+    # If using a single cached reference for a batch of RGB images, repeat it across batch.
+    if llava_input.shape[0] == 1 and batch_size > 1:
+        llava_input = llava_input.expand(batch_size, *llava_input.shape[1:])
+    if use_autocast:
+        llava_input = llava_input.to(dtype=torch.bfloat16)
+    
+    autocast_ctx = (
+        torch.autocast(device_type=device.type, dtype=torch.bfloat16) 
+        if use_autocast else nullcontext()
+    )
+    with autocast_ctx:
+        llava_tokens = adapter(llava_input)
+    llava_tokens = llava_tokens.to(device=device, dtype=next(unet.parameters()).dtype)
+    # Ensure token batch matches RGB batch (some adapters may return batch=1 even if input was expanded).
+    if llava_tokens.shape[0] == 1 and batch_size > 1:
+        llava_tokens = llava_tokens.expand(batch_size, *llava_tokens.shape[1:])
+    
+    # Create null tokens for CFG
+    null_tokens = torch.zeros_like(llava_tokens)
+    
+    # Initialize random noise
+    latent_shape = (batch_size, 4, rgb_latents.shape[2], rgb_latents.shape[3])
+    latents = torch.randn(latent_shape, device=device, dtype=rgb_latents.dtype)
+    
+    # Setup scheduler
+    scheduler.set_timesteps(num_steps, device=device)
+    
+    zeros_rgb = torch.zeros_like(rgb_latents)
+    
+    for t in scheduler.timesteps:
+        t_batch = torch.full((batch_size,), int(t), device=device, dtype=torch.long)
+        
+        # 3 forward passes for dual-CFG
+        x8_full = torch.cat([latents, rgb_latents], dim=1)
+        x8_no_img = torch.cat([latents, zeros_rgb], dim=1)
+        
+        # Full conditioning
+        eps_full = unet(x8_full, t_batch, encoder_hidden_states=llava_tokens).sample
+        
+        # No text (image only)
+        eps_no_text = unet(x8_full, t_batch, encoder_hidden_states=null_tokens).sample
+        
+        # No image (text only)
+        eps_no_img = unet(x8_no_img, t_batch, encoder_hidden_states=llava_tokens).sample
+        
+        # Combine with dual-CFG
+        eps = eps_full + cfg_text * (eps_full - eps_no_text) + cfg_image * (eps_full - eps_no_img)
+        eps = torch.clamp(eps, -5.0, 5.0)  # Safety guard
+        
+        latents = scheduler.step(eps, t, latents, eta=0.0).prev_sample
+    
+    # Decode
+    pred_tir = vae.decode(latents / vae.config.scaling_factor).sample
+    pred_tir = (pred_tir / 2 + 0.5).clamp(0, 1)
+    
+    return pred_tir
+
+
+def main():
+    args = parse_args()
+    validate_args(args)
+    
+    # Set seed
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+    
+    print("\n" + "="*80)
+    print(f"Example-Guided Translation - Mode: {args.mode.upper()}")
+    print("="*80)
+    print(f"Input image (B): {args.input_image}")
+    if args.mode == "two-image":
+        print(f"Reference image (A): {args.reference_image}")
+    else:
+        print(f"Reference cache: {args.reference_cache}")
+    print(f"Output: {args.output}")
+    print(f"Checkpoint: {args.checkpoint}")
+    print(f"DDIM steps: {args.num_steps}")
+    print(f"CFG scales: text={args.cfg_text}, image={args.cfg_image}")
+    print("="*80)
+    
+    # Load models
+    need_llava = (args.mode == "two-image")
+    vae, unet, adapter, scheduler, llava_extractor, device = load_models(args, need_llava=need_llava)
+    
+    # Get reference hidden states
+    print("\n" + "-"*80)
+    if args.mode == "two-image":
+        print(f"Extracting features from reference image: {args.reference_image}")
+        ref_tensor, ref_pil = load_and_prepare_image(
+            args.reference_image,
+            target_size=args.target_size,
+            device=device
+        )
+        llava_hidden = llava_extractor.extract_hidden_states([ref_pil], [args.llava_prompt])
+        print(f"✓ Extracted hidden states: {llava_hidden.shape}")
+    else:
+        print(f"Loading cached reference: {args.reference_cache}")
+        llava_hidden = load_llava_cache(Path(args.reference_cache), device)
+        print(f"✓ Loaded hidden states: {llava_hidden.shape}")
+    print("-"*80)
+    
+    # Load input image (Image B)
+    print(f"\nLoading input image: {args.input_image}")
+    input_tensor, input_pil = load_and_prepare_image(
+        args.input_image,
+        target_size=args.target_size,
+        device=device
+    )
+    print(f"✓ Input image loaded: {input_tensor.shape}")
+    
+    # Translate
+    print("\nTranslating image...")
+    pred_tir = translate_image(
+        vae, unet, adapter, scheduler,
+        input_tensor, llava_hidden,
+        num_steps=args.num_steps,
+        cfg_text=args.cfg_text,
+        cfg_image=args.cfg_image,
+        device=device
+    )
+    
+    # Save output
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    save_image(pred_tir.cpu(), str(output_path))
+    
+    print("="*80)
+    print(f"✓ Translation complete!")
+    print(f"✓ Saved to: {output_path}")
+    print("="*80)
+
+
+if __name__ == "__main__":
+    main()
+

infer_palette.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# Run TherA inference with weather/style palette reference caches.
+#
+# Usage:
+#   bash scripts/infer_palette.sh /path/to/rgb [output_dir]
+#
+# Requires weights/reference_caches/{SUNNY,CLOUDY,RAINY,NIGHT}.pt
+
+set -euo pipefail
+
+cd "$(dirname "$0")/.."
+
+RGB_DIR="${1:?RGB directory required}"
+OUT_BASE="${2:-preds_palette}"
+
+for PAL in SUNNY CLOUDY RAINY NIGHT; do
+  CACHE="weights/reference_caches/${PAL}.pt"
+  if [[ ! -f "$CACHE" ]]; then
+    echo "Skipping ${PAL}: missing ${CACHE}"
+    continue
+  fi
+  echo ">>> ${PAL}"
+  python infer_custom.py \
+    --rgb-dir "$RGB_DIR" \
+    --output-dir "${OUT_BASE}/${PAL}" \
+    --reference-cache "$CACHE"
+done
+
+echo ""
+echo "Done. Outputs in: ${OUT_BASE}/"

lavi_ip2p/__init__.py

@@ -0,0 +1,8 @@
+"""
+LaVi-IP2P Module
+
+RGB→TIR translation using LaVi-Bridge + cached LLaVA semantics.
+"""
+
+__version__ = "0.1.0"
+

lavi_ip2p/llava_adapter.py

@@ -0,0 +1,135 @@
+"""
+LLaVA Token Adapter for LaVi-IP2P
+
+Adapts cached LLaVA hidden states (4096-dim) to LaVi-Bridge UNet cross-attention (768-dim).
+Uses LaVi-Bridge's TextAdapter with optional preprocessing (RMSNorm, scaling).
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional
+import sys
+import os
+
+# Add LaVi-Bridge to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../LaVi-Bridge"))
+from modules.adapters import TextAdapter
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+    
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    
+    def forward(self, x):
+        norm = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return x * norm * self.weight
+
+
+class LLaVAAdapter(nn.Module):
+    """
+    Adapter: LLaVA hidden (4096) → LaVi-Bridge UNet conditioning (768).
+    
+    FROZEN by default (use pretrained LaVi-Bridge TextAdapter as-is).
+    """
+    
+    def __init__(
+        self,
+        adapter_path: str,
+        use_rms_norm: bool = False,
+        learnable_scale: bool = False,
+        init_scale: float = 1.0,
+        freeze_adapter: bool = True,
+    ):
+        super().__init__()
+        
+        self.use_rms_norm = use_rms_norm
+        self.learnable_scale = learnable_scale
+        
+        # Load pretrained TextAdapter
+        print(f"\nLoading TextAdapter from {adapter_path}...")
+        self.adapter = TextAdapter.from_pretrained(adapter_path)
+        print(f"✓ TextAdapter loaded (4096→768)")
+        
+        # Freeze by default
+        if freeze_adapter:
+            for param in self.adapter.parameters():
+                param.requires_grad = False
+            print("  ✓ FROZEN (using pretrained weights)")
+        else:
+            print("  ⚠️  TRAINABLE")
+        
+        # Optional: RMSNorm
+        if use_rms_norm:
+            self.rms_norm = RMSNorm(4096, eps=1e-6)
+            print("✓ RMSNorm enabled")
+        
+        # Optional: Learnable scale
+        if learnable_scale:
+            self.scale = nn.Parameter(torch.tensor(init_scale))
+            print(f"✓ Learnable scale (init={init_scale})")
+        else:
+            self.register_buffer('scale', torch.tensor(init_scale))
+    
+    def forward(self, llava_hidden: torch.Tensor) -> torch.Tensor:
+        """LLaVA (B,L,4096) → UNet conditioning (B,L,768)"""
+        x = llava_hidden
+        
+        if self.use_rms_norm:
+            x = self.rms_norm(x)
+        
+        x = self.adapter(x).sample
+        
+        if self.learnable_scale:
+            x = x * self.scale
+        
+        return x
+    
+    def get_stats(self):
+        """Get adapter statistics for logging"""
+        stats = {}
+        if self.learnable_scale:
+            stats['scale'] = self.scale.item()
+        return stats
+
+
+def create_llava_adapter(
+    adapter_path: str,
+    use_rms_norm: bool = False,
+    learnable_scale: bool = False,
+    init_scale: float = 1.0,
+    freeze_adapter: bool = True,
+) -> LLaVAAdapter:
+    """
+    Factory: Create LLaVA adapter.
+    
+    DEFAULT (simple mode):
+      - Freeze TextAdapter (use LaVi-Bridge pretrained)
+      - No RMSNorm, no scale
+      - Train UNet only
+    """
+    
+    print("\n" + "="*60)
+    print("Creating LLaVA Token Adapter")
+    print("="*60)
+    print(f"Mode: {'SIMPLE (frozen)' if freeze_adapter else 'ADVANCED (trainable)'}")
+    
+    adapter = LLaVAAdapter(
+        adapter_path=adapter_path,
+        use_rms_norm=use_rms_norm,
+        learnable_scale=learnable_scale,
+        init_scale=init_scale,
+        freeze_adapter=freeze_adapter,
+    )
+    
+    total = sum(p.numel() for p in adapter.parameters())
+    trainable = sum(p.numel() for p in adapter.parameters() if p.requires_grad)
+    
+    print(f"\nParameters: {total:,} total, {trainable:,} trainable")
+    print("="*60 + "\n")
+    
+    return adapter
+

lavi_ip2p/unet_8ch.py

@@ -0,0 +1,161 @@
+"""
+UNet 8-channel modification for LaVi-IP2P
+
+Safely converts LaVi-Bridge UNet from 4-channel to 8-channel input
+by expanding conv_in with zero-initialized extra channels (InstructPix2Pix style).
+"""
+
+import torch
+import torch.nn as nn
+from diffusers import UNet2DConditionModel
+from typing import Optional
+
+
+def convert_unet_to_8ch(
+    unet: UNet2DConditionModel,
+    init_scale: float = 1e-5,
+) -> UNet2DConditionModel:
+    """
+    Convert UNet from 4-channel to 8-channel input.
+    
+    Strategy (InstructPix2Pix style):
+    1. Copy first 4 channels from pretrained conv_in weights
+    2. Initialize extra 4 channels with small random values (scaled by init_scale)
+    3. This ensures the UNet starts close to the original behavior
+    
+    Args:
+        unet: Pretrained 4-channel UNet
+        init_scale: Scale for initializing extra channels (default: 1e-5 for stability)
+    
+    Returns:
+        Modified 8-channel UNet
+    """
+    
+    if unet.config.in_channels == 8:
+        print("⚠️  UNet already has 8 input channels, skipping conversion")
+        return unet
+    
+    if unet.config.in_channels != 4:
+        raise ValueError(f"Expected 4-channel UNet, got {unet.config.in_channels}")
+    
+    print("\n" + "="*60)
+    print("Converting UNet to 8-channel input (IP2P style)")
+    print("="*60)
+    
+    old_conv = unet.conv_in
+    print(f"Original conv_in: {old_conv}")
+    print(f"  in_channels: {old_conv.in_channels}")
+    print(f"  out_channels: {old_conv.out_channels}")
+    print(f"  kernel_size: {old_conv.kernel_size}")
+    print(f"  weight shape: {old_conv.weight.shape}")
+    
+    # Create new 8-channel conv
+    new_conv = nn.Conv2d(
+        in_channels=8,
+        out_channels=old_conv.out_channels,
+        kernel_size=old_conv.kernel_size,
+        stride=old_conv.stride,
+        padding=old_conv.padding,
+        bias=(old_conv.bias is not None),
+    )
+    
+    # Initialize weights
+    with torch.no_grad():
+        # Copy first 4 channels from pretrained weights
+        new_conv.weight[:, :4, :, :].copy_(old_conv.weight)
+        
+        # Initialize extra 4 channels with small random values
+        # This ensures the UNet starts close to original behavior
+        nn.init.normal_(new_conv.weight[:, 4:, :, :], mean=0.0, std=init_scale)
+        
+        # Copy bias if present
+        if old_conv.bias is not None:
+            new_conv.bias.copy_(old_conv.bias)
+    
+    # Replace conv_in
+    unet.conv_in = new_conv
+    
+    # Update config
+    unet.register_to_config(in_channels=8)
+    
+    print(f"\n✓ New conv_in: {new_conv}")
+    print(f"  weight shape: {new_conv.weight.shape}")
+    print(f"  Channels 0-3: copied from pretrained")
+    print(f"  Channels 4-7: initialized with std={init_scale}")
+    print("="*60 + "\n")
+    
+    return unet
+
+
+def load_lavi_unet_8ch(
+    pretrained_path: str,
+    lora_path: Optional[str] = None,
+    lora_rank: int = 32,
+    init_scale: float = 1e-5,
+    device: str = "cuda",
+) -> UNet2DConditionModel:
+    """
+    Load LaVi-Bridge UNet with LoRA and convert to 8-channel.
+    
+    Args:
+        pretrained_path: Path to pretrained SD UNet (e.g., "CompVis/stable-diffusion-v1-4")
+        lora_path: Path to LaVi-Bridge LoRA weights (lora_vis.pt)
+        lora_rank: LoRA rank (default: 32)
+        init_scale: Scale for extra channels
+        device: Device to load on
+    
+    Returns:
+        8-channel UNet with LaVi-Bridge LoRA applied
+    """
+    
+    print("\n" + "="*60)
+    print("Loading LaVi-Bridge UNet (8-channel)")
+    print("="*60)
+    
+    # Load base UNet
+    print(f"\n[1/3] Loading base UNet from {pretrained_path}...")
+    
+    # Determine subfolder based on filesystem check
+    from pathlib import Path
+    subfolder = "unet" if Path(pretrained_path).is_dir() and (Path(pretrained_path) / "unet").exists() else None
+    
+    unet = UNet2DConditionModel.from_pretrained(
+        pretrained_path,
+        subfolder=subfolder,
+    )
+    print(f"✓ Base UNet loaded (subfolder: {subfolder})")
+    
+    # Apply LoRA if provided
+    if lora_path is not None:
+        print(f"\n[2/3] Applying LaVi-Bridge LoRA from {lora_path}...")
+        
+        # Import LoRA utilities (robust path resolution)
+        import sys
+        import os
+        lavi_bridge_path = os.path.join(os.path.dirname(__file__), "../LaVi-Bridge")
+        sys.path.insert(0, lavi_bridge_path)
+        from modules.lora import monkeypatch_or_replace_lora_extended
+        
+        VIS_REPLACE_MODULES = {"ResnetBlock2D", "CrossAttention", "Attention", "GEGLU"}
+        
+        monkeypatch_or_replace_lora_extended(
+            unet,
+            torch.load(lora_path, map_location='cpu'),
+            r=lora_rank,
+            target_replace_module=VIS_REPLACE_MODULES,
+        )
+        print(f"✓ LoRA applied (rank={lora_rank})")
+    else:
+        print(f"\n[2/3] No LoRA path provided, using base UNet")
+    
+    # Convert to 8-channel
+    print(f"\n[3/3] Converting to 8-channel input...")
+    unet = convert_unet_to_8ch(unet, init_scale=init_scale)
+    
+    # Move to device
+    unet = unet.to(device)
+    print(f"✓ Moved to {device}")
+    print("="*60 + "\n")
+    
+    return unet
+

llava/__init__.py

@@ -0,0 +1 @@
+from .model import LlavaLlamaForCausalLM

llava/constants.py

@@ -0,0 +1,16 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+IMAGE_PLACEHOLDER = "<image-placeholder>"
+
+# IMG tokens for VLM-guided diffusion (16 tokens for edit conditioning)
+IMG_TOKENS = [f"[IMG{i}]" for i in range(16)]

llava/conversation.py

@@ -0,0 +1,396 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n" if len(msg) > 0 else msg
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def process_image(self, image, image_process_mode, return_pil=False, image_format='PNG', max_len=1344, min_len=672):
+        if image_process_mode == "Pad":
+            def expand2square(pil_img, background_color=(122, 116, 104)):
+                width, height = pil_img.size
+                if width == height:
+                    return pil_img
+                elif width > height:
+                    result = Image.new(pil_img.mode, (width, width), background_color)
+                    result.paste(pil_img, (0, (width - height) // 2))
+                    return result
+                else:
+                    result = Image.new(pil_img.mode, (height, height), background_color)
+                    result.paste(pil_img, ((height - width) // 2, 0))
+                    return result
+            image = expand2square(image)
+        elif image_process_mode in ["Default", "Crop"]:
+            pass
+        elif image_process_mode == "Resize":
+            image = image.resize((336, 336))
+        else:
+            raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+        if max(image.size) > max_len:
+            max_hw, min_hw = max(image.size), min(image.size)
+            aspect_ratio = max_hw / min_hw
+            shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+            longest_edge = int(shortest_edge * aspect_ratio)
+            W, H = image.size
+            if H > W:
+                H, W = longest_edge, shortest_edge
+            else:
+                H, W = shortest_edge, longest_edge
+            image = image.resize((W, H))
+        if return_pil:
+            return image
+        else:
+            buffered = BytesIO()
+            image.save(buffered, format=image_format)
+            img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+            return img_b64_str
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    image = self.process_image(image, image_process_mode, return_pil=return_pil)
+                    images.append(image)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    img_b64_str = self.process_image(
+                        image, "Default", return_pil=False,
+                        image_format='JPEG')
+                    img_str = f'<img src="data:image/jpeg;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+conv_mistral_instruct = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="",
+    sep2="</s>",
+)
+
+conv_chatml_direct = Conversation(
+    system="""<|im_start|>system
+Answer the questions.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+default_conversation = conv_vicuna_v1
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "mistral_instruct": conv_mistral_instruct,
+    "chatml_direct": conv_chatml_direct,
+    "mistral_direct": conv_chatml_direct,
+
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

llava/mm_utils.py

@@ -0,0 +1,247 @@
+from PIL import Image
+from io import BytesIO
+import base64
+import torch
+import math
+import ast
+
+from transformers import StoppingCriteria
+from llava.constants import IMAGE_TOKEN_INDEX
+
+
+def select_best_resolution(original_size, possible_resolutions):
+    """
+    Selects the best resolution from a list of possible resolutions based on the original size.
+
+    Args:
+        original_size (tuple): The original size of the image in the format (width, height).
+        possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
+
+    Returns:
+        tuple: The best fit resolution in the format (width, height).
+    """
+    original_width, original_height = original_size
+    best_fit = None
+    max_effective_resolution = 0
+    min_wasted_resolution = float('inf')
+
+    for width, height in possible_resolutions:
+        scale = min(width / original_width, height / original_height)
+        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)
+        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
+        wasted_resolution = (width * height) - effective_resolution
+
+        if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
+            max_effective_resolution = effective_resolution
+            min_wasted_resolution = wasted_resolution
+            best_fit = (width, height)
+
+    return best_fit
+
+
+def resize_and_pad_image(image, target_resolution):
+    """
+    Resize and pad an image to a target resolution while maintaining aspect ratio.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        target_resolution (tuple): The target resolution (width, height) of the image.
+
+    Returns:
+        PIL.Image.Image: The resized and padded image.
+    """
+    original_width, original_height = image.size
+    target_width, target_height = target_resolution
+
+    scale_w = target_width / original_width
+    scale_h = target_height / original_height
+
+    if scale_w < scale_h:
+        new_width = target_width
+        new_height = min(math.ceil(original_height * scale_w), target_height)
+    else:
+        new_height = target_height
+        new_width = min(math.ceil(original_width * scale_h), target_width)
+
+    # Resize the image
+    resized_image = image.resize((new_width, new_height))
+
+    new_image = Image.new('RGB', (target_width, target_height), (0, 0, 0))
+    paste_x = (target_width - new_width) // 2
+    paste_y = (target_height - new_height) // 2
+    new_image.paste(resized_image, (paste_x, paste_y))
+
+    return new_image
+
+
+def divide_to_patches(image, patch_size):
+    """
+    Divides an image into patches of a specified size.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        patch_size (int): The size of each patch.
+
+    Returns:
+        list: A list of PIL.Image.Image objects representing the patches.
+    """
+    patches = []
+    width, height = image.size
+    for i in range(0, height, patch_size):
+        for j in range(0, width, patch_size):
+            box = (j, i, j + patch_size, i + patch_size)
+            patch = image.crop(box)
+            patches.append(patch)
+
+    return patches
+
+
+def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
+    """
+    Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
+
+    Args:
+        image_size (tuple): The size of the input image in the format (width, height).
+        grid_pinpoints (str): A string representation of a list of possible resolutions.
+        patch_size (int): The size of each image patch.
+
+    Returns:
+        tuple: The shape of the image patch grid in the format (width, height).
+    """
+    if type(grid_pinpoints) is list:
+        possible_resolutions = grid_pinpoints
+    else:
+        possible_resolutions = ast.literal_eval(grid_pinpoints)
+    width, height = select_best_resolution(image_size, possible_resolutions)
+    return width // patch_size, height // patch_size
+
+
+def process_anyres_image(image, processor, grid_pinpoints):
+    """
+    Process an image with variable resolutions.
+
+    Args:
+        image (PIL.Image.Image): The input image to be processed.
+        processor: The image processor object.
+        grid_pinpoints (str): A string representation of a list of possible resolutions.
+
+    Returns:
+        torch.Tensor: A tensor containing the processed image patches.
+    """
+    if type(grid_pinpoints) is list:
+        possible_resolutions = grid_pinpoints
+    else:
+        possible_resolutions = ast.literal_eval(grid_pinpoints)
+    best_resolution = select_best_resolution(image.size, possible_resolutions)
+    image_padded = resize_and_pad_image(image, best_resolution)
+
+    patches = divide_to_patches(image_padded, processor.crop_size['height'])
+
+    image_original_resize = image.resize((processor.size['shortest_edge'], processor.size['shortest_edge']))
+
+    image_patches = [image_original_resize] + patches
+    image_patches = [processor.preprocess(image_patch, return_tensors='pt')['pixel_values'][0]
+                     for image_patch in image_patches]
+    return torch.stack(image_patches, dim=0)
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg):
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    new_images = []
+    if image_aspect_ratio == 'pad':
+        for image in images:
+            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            new_images.append(image)
+    elif image_aspect_ratio == "anyres":
+        for image in images:
+            image = process_anyres_image(image, image_processor, model_cfg.image_grid_pinpoints)
+            new_images.append(image)
+    else:
+        return image_processor(images, return_tensors='pt')['pixel_values']
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    offset = 0
+    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+
+    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+        input_ids.extend(x[offset:])
+
+    if return_tensors is not None:
+        if return_tensors == 'pt':
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f'Unsupported tensor type: {return_tensors}')
+    return input_ids
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith('checkpoint-'):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+    
+    def call_for_batch(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            truncated_output_ids = output_ids[0, -keyword_id.shape[0]:]
+            if torch.equal(truncated_output_ids, keyword_id):
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False
+    
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        outputs = []
+        for i in range(output_ids.shape[0]):
+            outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+        return all(outputs)

llava/rgbt_dataset.py

@@ -0,0 +1,328 @@
+# llava/data/rgb_tir_dataset.py
+import os
+import copy
+import json
+from typing import Dict, Sequence, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms as tvt
+
+from llava.constants import (
+    IGNORE_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+)
+from llava import conversation as conversation_lib
+from llava.mm_utils import tokenizer_image_token
+
+def _safe_replace_rgb_with_tir(path: str) -> str:
+    # handle /RGB/ or \RGB\ on any OS
+    path2 = path.replace(os.sep + "RGB" + os.sep, os.sep + "TIR" + os.sep)
+    path2 = path2.replace("/RGB/", "/TIR/").replace("\\RGB\\", "\\TIR\\")
+    return path2
+
+def _expand2square(pil_img: Image.Image, background_color=(0, 0, 0)):
+    w, h = pil_img.size
+    if w == h:
+        return pil_img
+    if w > h:
+        result = Image.new(pil_img.mode, (w, w), background_color)
+        result.paste(pil_img, (0, (w - h) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (h, h), background_color)
+        result.paste(pil_img, ((h - w) // 2, 0))
+        return result
+
+class RGBTIRSupervisedDataset(Dataset):
+    """
+    LLaVA-compatible dataset that also (optionally) returns RGB/TIR pairs
+    for diffusion training (instructpix2pix). When diffusion pairs are off,
+    it behaves like the stock LazySupervisedDataset.
+
+    Keys:
+      - Always for LLaVA: input_ids, labels, image (-> collated as `images`)
+      - Optional for diffusion: rgb_diffusion, tir_diffusion
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer,
+        image_processor,
+        image_folder: Optional[str] = None,
+        image_aspect_ratio: str = "square",  # {"square","pad"}
+        model_max_length: int = 2048,
+        # diffusion extras
+        return_diffusion_pairs: bool = False,
+        diffusion_size: int = 512,
+        tir_as_rgb: bool = True,  # True: convert TIR to 3ch; False: keep 1ch ("L")
+        mm_use_im_start_end: bool = False,   # NEW
+        warn_missing_tir: bool = False,      # optional
+    ):
+        super().__init__()
+        self.data = json.load(open(data_path, "r"))
+        self.tokenizer = tokenizer
+        self.image_processor = image_processor
+        self.image_folder = image_folder or ""
+        self.image_aspect_ratio = image_aspect_ratio
+        self.model_max_length = model_max_length
+        self.return_diffusion_pairs = return_diffusion_pairs
+        self.diffusion_size = diffusion_size
+        self.tir_as_rgb = tir_as_rgb
+        self.mm_use_im_start_end = mm_use_im_start_end   # NEW
+        self.warn_missing_tir = warn_missing_tir
+
+
+        def make_rgb_transform():
+            return tvt.Compose([
+                tvt.CenterCrop(min(self.diffusion_size, 10**9)),  # overwritten next line
+                tvt.Resize(self.diffusion_size, interpolation=Image.BICUBIC),
+                tvt.ToTensor(),
+                tvt.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+            ])
+
+        # CenterCrop to exact square first, then Resize
+        def square_then_resize(mode_rgb: bool):
+            norms = ([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) if mode_rgb else ([0.5], [0.5])
+            return tvt.Compose([
+                # crop to min side square
+                tvt.Lambda(lambda im: _expand2square(im, (0, 0, 0)) if im.size[0]!=im.size[1] else im),
+                tvt.Resize(self.diffusion_size, interpolation=Image.BICUBIC),
+                tvt.ToTensor(),
+                tvt.Normalize(*norms),
+            ])
+
+        self.diffusion_transform_rgb = square_then_resize(True)
+        self.diffusion_transform_tir = square_then_resize(self.tir_as_rgb)
+        # diffusion transforms ([-1,1])
+        # self.diffusion_transform_rgb = tvt.Compose([
+        #     tvt.Resize(self.diffusion_size, interpolation=Image.BICUBIC),
+        #     tvt.CenterCrop(self.diffusion_size),
+        #     tvt.ToTensor(),
+        #     tvt.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        # ])
+        # allow 1ch or 3ch for TIR
+        # if self.tir_as_rgb:
+        #     self.diffusion_transform_tir = tvt.Compose([
+        #         tvt.Resize(self.diffusion_size, interpolation=Image.BICUBIC),
+        #         tvt.CenterCrop(self.diffusion_size),
+        #         tvt.ToTensor(),
+        #         tvt.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        #     ])
+        # else:
+        #     self.diffusion_transform_tir = tvt.Compose([
+        #         tvt.Resize(self.diffusion_size, interpolation=Image.BICUBIC),
+        #         tvt.CenterCrop(self.diffusion_size),
+        #         tvt.ToTensor(),                     # 1 x H x W
+        #         tvt.Normalize([0.5], [0.5]),
+        #     ])
+
+        # in RGBTIRSupervisedDataset.__init__(...)
+        self.inject_edit_tokens = True   # optional flag
+        self.edit_span = " ".join([f"<edit_{k}>" for k in range(8)])
+
+
+    def __len__(self):
+        return len(self.data)
+
+
+    @property
+    def lengths(self):
+        length_list = []
+        for sample in self.data:
+            img_tokens = 128 if 'image' in sample else 0
+            length_list.append(sum(len(conv['value'].split()) for conv in sample['conversations']) + img_tokens)
+        return length_list
+
+    @property
+    def modality_lengths(self):
+        length_list = []
+        for sample in self.data:
+            cur_len = sum(len(conv['value'].split()) for conv in sample['conversations'])
+            cur_len = cur_len if 'image' in sample else -cur_len
+            length_list.append(cur_len)
+        return length_list
+
+    # add this helper in the class
+    def _inject_edit_tokens_into_conversation(self, conv_rounds: list) -> list:
+        """
+        Ensure the last assistant ('gpt') turn is exactly the 8 edit tokens.
+        If there is no assistant turn yet, append one.
+        conv_rounds: [{'from': 'human'|'gpt', 'value': str}, ...]
+        """
+        if not self.inject_edit_tokens:
+            return conv_rounds
+
+        # find last assistant turn
+        last_gpt = None
+        for i in reversed(range(len(conv_rounds))):
+            if conv_rounds[i].get("from") == "gpt":
+                last_gpt = i
+                break
+
+        if last_gpt is None:
+            # append a new assistant turn with just the edit tokens
+            conv_rounds = conv_rounds + [{"from": "gpt", "value": self.edit_span}]
+        else:
+            # replace its content with exactly the edit tokens
+            conv_rounds = conv_rounds.copy()
+            conv_rounds[last_gpt] = {"from": "gpt", "value": self.edit_span}
+
+        return conv_rounds
+
+
+    
+    # --- helpers mirroring your current preprocess path ---
+    def _preprocess_multimodal(self, sources):
+        use_im_tokens = self.mm_use_im_start_end
+        for source in sources:
+            for sentence in source:
+                if DEFAULT_IMAGE_TOKEN in sentence["value"]:
+                    sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, "").strip()
+                    sentence["value"] = DEFAULT_IMAGE_TOKEN + "\n" + sentence["value"]
+                    sentence["value"] = sentence["value"].strip()
+                replace_token = DEFAULT_IMAGE_TOKEN
+                if use_im_tokens:
+                    replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
+                sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, replace_token)
+        return sources
+
+
+    def _preprocess_text(self, sources, has_image: bool):
+        # Follow the same path as llava.data.preprocess_v1 / llama_2 depending on conv template
+        conv = conversation_lib.default_conversation.copy()
+        roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
+        conversations = []
+
+        for i, source in enumerate(sources):
+            if roles[source[0]["from"]] != conv.roles[0]:
+                source = source[1:]  # skip non-human first
+            conv.messages = []
+            for j, sentence in enumerate(source):
+                role = roles[sentence["from"]]
+                assert role == conv.roles[j % 2]
+                conv.append_message(role, sentence["value"])
+            conversations.append(conv.get_prompt())
+
+        if has_image:
+            input_ids = torch.stack(
+                [tokenizer_image_token(p, self.tokenizer, return_tensors="pt") for p in conversations], dim=0
+            )
+        else:
+            input_ids = self.tokenizer(
+                conversations, return_tensors="pt", padding="longest",
+                max_length=self.tokenizer.model_max_length, truncation=True
+            ).input_ids
+
+        targets = input_ids.clone()
+
+        # mask human parts (match llava.data.preprocess_v1 / llama_2 implementations)
+        if conv.sep_style == conversation_lib.SeparatorStyle.TWO:
+            sep = conv.sep + conv.roles[1] + ": "
+            for conversation, target in zip(conversations, targets):
+                total_len = int(target.ne(self.tokenizer.pad_token_id).sum())
+                rounds = conversation.split(conv.sep2)
+                cur_len = 1
+                target[:cur_len] = IGNORE_INDEX
+                for i, rou in enumerate(rounds):
+                    if rou == "":
+                        break
+                    parts = rou.split(sep)
+                    if len(parts) != 2:
+                        break
+                    parts[0] += sep
+                    if has_image:
+                        round_len = len(tokenizer_image_token(rou, self.tokenizer))
+                        instruction_len = len(tokenizer_image_token(parts[0], self.tokenizer)) - 2
+                    else:
+                        round_len = len(self.tokenizer(rou).input_ids)
+                        instruction_len = len(self.tokenizer(parts[0]).input_ids) - 2
+                    target[cur_len: cur_len + instruction_len] = IGNORE_INDEX
+                    cur_len += round_len
+                target[cur_len:] = IGNORE_INDEX
+                if cur_len != total_len:
+                    target[:] = IGNORE_INDEX  # safe path
+        else:
+            # fall back to generic masking path (rare given your templates)
+            target = targets  # already cloned
+        return dict(input_ids=input_ids[0], labels=targets[0])
+
+    def _load_rgb_for_llava(self, pil_img: Image.Image):
+        # honor aspect ratio option used by your current code
+        if self.image_aspect_ratio == "pad":
+            bg = tuple(int(x * 255) for x in self.image_processor.image_mean)
+            pil_img = _expand2square(pil_img, bg)
+        # LLaVA’s processor returns dict(pixel_values=...)
+        return self.image_processor.preprocess(pil_img, return_tensors="pt")["pixel_values"][0]
+
+    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
+        sample = self.data[idx]
+        rgb_rel = sample.get("image")
+        assert rgb_rel is not None, "Missing 'image' in sample"
+
+        rgb_path = os.path.join(self.image_folder, rgb_rel)
+        rgb = Image.open(rgb_path).convert("RGB")
+
+        # paired TIR (optional)
+        tir_path = _safe_replace_rgb_with_tir(rgb_path)
+        tir_exists = os.path.exists(tir_path)
+        if self.return_diffusion_pairs and (not tir_exists) and self.warn_missing_tir:
+            print(f"[RGBTIRDataset] Missing TIR for: {rgb_path}")
+
+        # LLaVA image tensor
+        image_tensor = self._load_rgb_for_llava(rgb)
+
+        # --------- NEW: build MGIE-style two-turn conversation from your JSON ---------
+        # human value is the "instruction"; gpt value is the "expressive" text
+        convs = sample["conversations"]
+        assert len(convs) >= 2, "Expected at least 2 turns (human, gpt)"
+
+        # take the first human turn as instruction
+        human_val = convs[0]["value"]
+        # strip a leading <image> (if present) to avoid double image tokens; we’ll insert ours below
+        if human_val.lstrip().startswith("<image>"):
+            human_val = human_val.lstrip()[len("<image>"):].lstrip()
+
+        # expressive text is the first gpt turn
+        expressive = convs[1]["value"].strip()
+
+        # choose the 8 tokens: use the ones you ALREADY added to the tokenizer/model
+        EDIT_TOKENS = [f"[IMG{k}]" for k in range(8)]
+        edit_span = " ".join(EDIT_TOKENS)
+
+        # build the two messages in MGIE style
+        query = f"{human_val}\n{DEFAULT_IMAGE_TOKEN}"            # human: instruction + <image>
+        answer = f"{expressive} {edit_span}"                     # gpt: expressive + 8 tokens
+
+        sources = [[
+            {"from": "human", "value": query},
+            {"from": "gpt",   "value": answer},
+        ]]
+
+        # make sure image token formatting matches your config
+        sources = self._preprocess_multimodal(copy.deepcopy(sources))
+        text_dict = self._preprocess_text(sources, has_image=True)
+
+        # pack outputs for the trainer
+        out = {
+            "input_ids": text_dict["input_ids"],
+            "labels": text_dict["labels"],
+            "image": image_tensor,
+            "image_sizes": [rgb.size],
+        }
+
+        # --------- optional diffusion pairs ---------
+        if self.return_diffusion_pairs:
+            rgb_for_diff = rgb
+            if tir_exists:
+                tir_for_diff = Image.open(tir_path).convert("RGB" if self.tir_as_rgb else "L")
+            else:
+                tir_for_diff = rgb_for_diff.convert("RGB" if self.tir_as_rgb else "L")
+
+            out["p2p_inp"] = self.diffusion_transform_rgb(rgb_for_diff)
+            out["p2p_ans"] = self.diffusion_transform_tir(tir_for_diff)
+
+        return out
+
+

llava/train/llama_flash_attn_monkey_patch.py

@@ -0,0 +1,115 @@
+from typing import Optional, Tuple
+import warnings
+
+import torch
+
+import transformers
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
+except ImportError:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+from flash_attn.bert_padding import unpad_input, pad_input
+
+
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    if output_attentions:
+        warnings.warn(
+            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+        )
+
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = (
+        self.q_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    key_states = (
+        self.k_proj(hidden_states)
+        .view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    value_states = (
+        self.v_proj(hidden_states)
+        .view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        .transpose(1, 2)
+    )  # shape: (b, num_heads, s, head_dim)
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+
+    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, position_ids
+    )
+
+    if past_key_value is not None:
+        # reuse k, v
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
+
+    # repeat k/v heads if n_kv_heads < n_heads
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+    # Transform the data into the format required by flash attention
+    qkv = torch.stack([query_states, key_states, value_states], dim=2)
+    qkv = qkv.transpose(1, 3)  # shape: [b, s, 3, num_heads, head_dim]
+    key_padding_mask = attention_mask
+
+    if key_padding_mask is None:
+        qkv = qkv.reshape(-1, 3, self.num_heads, self.head_dim)
+        cu_q_lens = torch.arange(
+            0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device
+        )
+        max_s = q_len
+        output = flash_attn_unpadded_qkvpacked_func(
+            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
+        )
+        output = output.view(bsz, q_len, -1)
+    else:
+        qkv = qkv.reshape(bsz, q_len, -1)
+        qkv, indices, cu_q_lens, max_s = unpad_input(qkv, key_padding_mask)
+        qkv = qkv.view(-1, 3, self.num_heads, self.head_dim)
+        output_unpad = flash_attn_unpadded_qkvpacked_func(
+            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
+        )
+        output_unpad = output_unpad.reshape(-1, self.num_heads * self.head_dim)
+        output = pad_input(output_unpad, indices, bsz, q_len)
+
+    return self.o_proj(output), None, past_key_value
+
+
+# Disable the transformation of the attention mask in LlamaModel as the flash attention
+# requires the attention mask to be the same as the key_padding_mask
+def _prepare_decoder_attention_mask(
+    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+):
+    # [bsz, seq_len]
+    return attention_mask
+
+
+def replace_llama_attn_with_flash_attn():
+    cuda_major, cuda_minor = torch.cuda.get_device_capability()
+    if cuda_major < 8:
+        warnings.warn(
+            "Flash attention is only supported on A100 or H100 GPU during training due to head dim > 64 backward."
+            "ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593"
+        )
+    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (
+        _prepare_decoder_attention_mask
+    )
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = forward

llava/train/llama_xformers_attn_monkey_patch.py

@@ -0,0 +1,129 @@
+"""
+Directly copied the code from https://raw.githubusercontent.com/oobabooga/text-generation-webui/main/modules/llama_attn_hijack.py and made some adjustments
+"""
+
+import logging
+import math
+from typing import Optional, Tuple
+
+import torch
+import transformers.models.llama.modeling_llama
+from torch import nn
+
+try:
+    import xformers.ops
+except ImportError:
+    logging.error("xformers not found! Please install it before trying to use it.")
+
+
+def replace_llama_attn_with_xformers_attn():
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = xformers_forward
+
+
+def xformers_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    # pylint: disable=duplicate-code
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = (
+        self.q_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    key_states = (
+        self.k_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    value_states = (
+        self.v_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+    (
+        query_states,
+        key_states,
+    ) = transformers.models.llama.modeling_llama.apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, position_ids
+    )
+    # [bsz, nh, t, hd]
+
+    if past_key_value is not None:
+        # reuse k, v, self_attention
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
+
+    # We only apply xformers optimizations if we don't need to output the whole attention matrix
+    if not output_attentions:
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        # This is a nasty hack. We know attention_mask in transformers is either LowerTriangular or all Zeros.
+        # We therefore check if one element in the upper triangular portion is zero. If it is, then the mask is all zeros.
+        if attention_mask is None or attention_mask[0, 0, 0, 1] == 0:
+            # input and output should be of form (bsz, q_len, num_heads, head_dim)
+            attn_output = xformers.ops.memory_efficient_attention(
+                query_states, key_states, value_states, attn_bias=None
+            )
+        else:
+            # input and output should be of form (bsz, q_len, num_heads, head_dim)
+            attn_output = xformers.ops.memory_efficient_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_bias=xformers.ops.LowerTriangularMask(),
+            )
+        attn_weights = None
+    else:
+        attn_weights = torch.matmul(
+            query_states, key_states.transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+            attn_weights = torch.max(
+                attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min)
+            )
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32
+        ).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+    attn_output = self.o_proj(attn_output)
+    return attn_output, attn_weights, past_key_value

llava/train/llava_trainer.py

@@ -0,0 +1,273 @@
+import os
+import torch
+import torch.nn as nn
+
+from torch.utils.data import Sampler
+
+from transformers import Trainer
+from transformers.trainer import (
+    is_sagemaker_mp_enabled,
+    get_parameter_names,
+    has_length,
+    ALL_LAYERNORM_LAYERS,
+    logger,
+)
+from typing import List, Optional
+
+
+def maybe_zero_3(param, ignore_status=False, name=None):
+    from deepspeed import zero
+    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+    if hasattr(param, "ds_id"):
+        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+            if not ignore_status:
+                print(name, 'no ignore status')
+        with zero.GatheredParameters([param]):
+            param = param.data.detach().cpu().clone()
+    else:
+        param = param.detach().cpu().clone()
+    return param
+
+
+def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):
+    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}
+    to_return = {k: maybe_zero_3(v, ignore_status=True, name=k).cpu() for k, v in to_return.items()}
+    return to_return
+
+
+def split_to_even_chunks(indices, lengths, num_chunks):
+    """
+    Split a list of indices into `chunks` chunks of roughly equal lengths.
+    """
+
+    if len(indices) % num_chunks != 0:
+        return [indices[i::num_chunks] for i in range(num_chunks)]
+
+    num_indices_per_chunk = len(indices) // num_chunks
+
+    chunks = [[] for _ in range(num_chunks)]
+    chunks_lengths = [0 for _ in range(num_chunks)]
+    for index in indices:
+        shortest_chunk = chunks_lengths.index(min(chunks_lengths))
+        chunks[shortest_chunk].append(index)
+        chunks_lengths[shortest_chunk] += lengths[index]
+        if len(chunks[shortest_chunk]) == num_indices_per_chunk:
+            chunks_lengths[shortest_chunk] = float("inf")
+
+    return chunks
+
+
+def get_modality_length_grouped_indices(lengths, batch_size, world_size, generator=None):
+    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
+    assert all(l != 0 for l in lengths), "Should not have zero length."
+    if all(l > 0 for l in lengths) or all(l < 0 for l in lengths):
+        # all samples are in the same modality
+        return get_length_grouped_indices(lengths, batch_size, world_size, generator=generator)
+    mm_indices, mm_lengths = zip(*[(i, l) for i, l in enumerate(lengths) if l > 0])
+    lang_indices, lang_lengths = zip(*[(i, -l) for i, l in enumerate(lengths) if l < 0])
+
+    mm_shuffle = [mm_indices[i] for i in get_length_grouped_indices(mm_lengths, batch_size, world_size, generator=None)]
+    lang_shuffle = [lang_indices[i] for i in get_length_grouped_indices(lang_lengths, batch_size, world_size, generator=None)]
+    megabatch_size = world_size * batch_size
+    mm_megabatches = [mm_shuffle[i : i + megabatch_size] for i in range(0, len(mm_shuffle), megabatch_size)]
+    lang_megabatches = [lang_shuffle[i : i + megabatch_size] for i in range(0, len(lang_shuffle), megabatch_size)]
+
+    last_mm = mm_megabatches[-1]
+    last_lang = lang_megabatches[-1]
+    additional_batch = last_mm + last_lang
+    megabatches = mm_megabatches[:-1] + lang_megabatches[:-1]
+    megabatch_indices = torch.randperm(len(megabatches), generator=generator)
+    megabatches = [megabatches[i] for i in megabatch_indices]
+
+    if len(additional_batch) > 0:
+        megabatches.append(sorted(additional_batch))
+
+    return [i for megabatch in megabatches for i in megabatch]
+
+
+def get_length_grouped_indices(lengths, batch_size, world_size, generator=None, merge=True):
+    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
+    indices = torch.randperm(len(lengths), generator=generator)
+    megabatch_size = world_size * batch_size
+    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]
+    megabatches = [split_to_even_chunks(megabatch, lengths, world_size) for megabatch in megabatches]
+
+    return [i for megabatch in megabatches for batch in megabatch for i in batch]
+
+
+class LengthGroupedSampler(Sampler):
+    r"""
+    Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
+    keeping a bit of randomness.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        world_size: int,
+        lengths: Optional[List[int]] = None,
+        generator=None,
+        group_by_modality: bool = False,
+    ):
+        if lengths is None:
+            raise ValueError("Lengths must be provided.")
+
+        self.batch_size = batch_size
+        self.world_size = world_size
+        self.lengths = lengths
+        self.generator = generator
+        self.group_by_modality = group_by_modality
+
+    def __len__(self):
+        return len(self.lengths)
+
+    def __iter__(self):
+        if self.group_by_modality:
+            indices = get_modality_length_grouped_indices(self.lengths, self.batch_size, self.world_size, generator=self.generator)
+        else:
+            indices = get_length_grouped_indices(self.lengths, self.batch_size, self.world_size, generator=self.generator)
+        return iter(indices)
+
+
+class LLaVATrainer(Trainer):
+
+    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
+        if self.train_dataset is None or not has_length(self.train_dataset):
+            return None
+
+        if self.args.group_by_modality_length:
+            lengths = self.train_dataset.modality_lengths
+            return LengthGroupedSampler(
+                self.args.train_batch_size,
+                world_size=self.args.world_size * self.args.gradient_accumulation_steps,
+                lengths=lengths,
+                group_by_modality=True,
+            )
+        else:
+            return super()._get_train_sampler()
+
+    def create_optimizer(self):
+        """
+        Setup the optimizer.
+
+        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
+        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
+        """
+        if is_sagemaker_mp_enabled():
+            return super().create_optimizer()
+
+        opt_model = self.model
+
+        if self.optimizer is None:
+            decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
+            decay_parameters = [name for name in decay_parameters if "bias" not in name]
+            if self.args.mm_projector_lr is not None:
+                projector_parameters = [name for name, _ in opt_model.named_parameters() if "mm_projector" in name]
+                optimizer_grouped_parameters = [
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in projector_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": self.args.weight_decay,
+                    },
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in projector_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": 0.0,
+                    },
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in projector_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": self.args.weight_decay,
+                        "lr": self.args.mm_projector_lr,
+                    },
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in projector_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": 0.0,
+                        "lr": self.args.mm_projector_lr,
+                    },
+                ]
+            else:
+                optimizer_grouped_parameters = [
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": self.args.weight_decay,
+                    },
+                    {
+                        "params": [
+                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
+                        ],
+                        "weight_decay": 0.0,
+                    },
+                ]
+
+            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
+
+            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+            if optimizer_cls.__name__ == "Adam8bit":
+                import bitsandbytes
+
+                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+
+                skipped = 0
+                for module in opt_model.modules():
+                    if isinstance(module, nn.Embedding):
+                        skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
+                        logger.info(f"skipped {module}: {skipped/2**20}M params")
+                        manager.register_module_override(module, "weight", {"optim_bits": 32})
+                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
+                logger.info(f"skipped: {skipped/2**20}M params")
+
+        return self.optimizer
+
+    def _save_checkpoint(self, model, trial, metrics=None):
+        if getattr(self.args, 'tune_mm_mlp_adapter', False):
+            from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
+
+            run_dir = self._get_output_dir(trial=trial)
+            output_dir = os.path.join(run_dir, checkpoint_folder)
+
+            # Only save Adapter
+            keys_to_match = ['mm_projector', 'vision_resampler']
+            if getattr(self.args, "use_im_start_end", False):
+                keys_to_match.extend(['embed_tokens', 'embed_in'])
+
+            weight_to_save = get_mm_adapter_state_maybe_zero_3(self.model.named_parameters(), keys_to_match)
+
+            if self.args.local_rank == 0 or self.args.local_rank == -1:
+                self.model.config.save_pretrained(output_dir)
+                torch.save(weight_to_save, os.path.join(output_dir, f'mm_projector.bin'))
+        else:
+            super(LLaVATrainer, self)._save_checkpoint(model, trial, metrics)
+
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        if getattr(self.args, 'tune_mm_mlp_adapter', False):
+            pass
+        else:
+            super(LLaVATrainer, self)._save(output_dir, state_dict)
+        
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs = model(**inputs)
+        loss = outputs.loss
+
+        logs = {}
+        # ALWAYS convert tensors -> floats before self.log(...)
+        if hasattr(outputs, "loss_ce") and outputs.loss_ce is not None:
+            logs["ce_loss"] = outputs.loss_ce.detach().float().mean().item()
+        if hasattr(outputs, "loss_edit") and outputs.loss_edit is not None:
+            logs["diffusion_loss"] = outputs.loss_edit.detach().float().mean().item()
+        # (optional) also log total loss as a scalar
+        logs["loss"] = loss.detach().float().item()
+
+        if logs:
+            self.log(logs)
+
+        return (loss, outputs) if return_outputs else loss

llava/train/train_vlm_ip2p_256.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""
+Training script for VLM-guided InstructPix2Pix (256x256 resolution).
+
+Optimized for faster training with 256x256 images.
+
+Key Features:
+1. 256x256 diffusion resolution (4x faster than 512x512)
+2. Choice of EditMapper: Simple MLP or MGIE-style Transformer
+3. All cross-attention blocks injected by default
+4. Compatible with base IP2P model (no need for fine-tuned checkpoint)
+
+Usage:
+    deepspeed --num_gpus=4 llava/train/train_vlm_ip2p_256.py \\
+        --llava-base-path checkpoints/llava-1.5-7b-hf \\
+        --llava-lora-path checkpoints/llava-miragehd-prior-bbox \\
+        --ip2p-pretrained ../diffusers/ip2p_ema_safetensor \\
+        --data-path data/llava_no_ranking_nobbox.json \\
+        --image-folder ../ \\
+        --output-dir ./outputs/ip_adapter_256_v1 \\
+        --batch-size 4 \\
+        --num-epochs 100 \\
+        --use-mgie-mapper \\
+        --deepspeed ./scripts/deepspeed_zero2_vlm_ip2p.json
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+try:
+    import deepspeed
+    DEEPSPEED_AVAILABLE = True
+except ImportError:
+    DEEPSPEED_AVAILABLE = False
+
+# Import the base trainer
+sys.path.append(str(Path(__file__).parent.parent.parent))
+from llava.train.train_vlm_guided_ip2p import VLMGuidedIP2PTrainer, parse_args
+
+
+class VLMGuidedIP2P256Trainer(VLMGuidedIP2PTrainer):
+    """
+    Trainer optimized for 256x256 resolution.
+    
+    Inherits from base trainer and overrides data loading
+    to use 256x256 images for diffusion.
+    """
+    
+    def setup_dataloaders(self):
+        """Setup dataloaders with 256x256 resolution for diffusion"""
+        from llava.data import create_dataloaders
+        
+        print("\n" + "="*80)
+        print("SETTING UP DATALOADERS (256x256 RESOLUTION)")
+        print("="*80)
+        
+        # Override diffusion_size to 256
+        original_diffusion_size = self.args.diffusion_size if hasattr(self.args, 'diffusion_size') else 512
+        self.args.diffusion_size = 256  # Force 256x256 for diffusion
+        
+        print(f"  Diffusion resolution: {self.args.diffusion_size}x{self.args.diffusion_size}")
+        print(f"  (Original was: {original_diffusion_size}x{original_diffusion_size})")
+        
+        # Call parent setup with modified args
+        super().setup_dataloaders()
+        
+        print("="*80 + "\n")
+
+
+def main():
+    # Parse arguments
+    args = parse_args()
+    
+    # Override defaults for 256x256 training
+    print("\n" + "="*80)
+    print("VLM-GUIDED IP2P TRAINING (256x256 OPTIMIZED)")
+    print("="*80)
+    print(f"Resolution: 256x256 (4x faster than 512x512)")
+    print(f"EditMapper: {'MGIE-style' if args.use_mgie_mapper else 'Simple MLP'}")
+    print(f"Target blocks: {'All cross-attention' if args.target_blocks is None else args.target_blocks}")
+    print(f"IP2P checkpoint: {args.ip2p_pretrained}")
+    
+    # Warn if using fine-tuned checkpoint
+    if args.ip2p_checkpoint:
+        print("\n⚠ WARNING: You specified --ip2p-checkpoint")
+        print("  For 256x256 training, you should use the base model (--ip2p-pretrained only)")
+        print("  Fine-tuned checkpoints may be trained on 512x512 resolution")
+        response = input("  Continue anyway? (y/N): ")
+        if response.lower() != 'y':
+            print("Exiting...")
+            return
+    
+    print("="*80 + "\n")
+    
+    # Create trainer
+    trainer = VLMGuidedIP2P256Trainer(args)
+    
+    # Train
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()
+

llava/utils.py

@@ -0,0 +1,126 @@
+import datetime
+import logging
+import logging.handlers
+import os
+import sys
+
+import requests
+
+from llava.constants import LOGDIR
+
+server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+handler = None
+
+
+def build_logger(logger_name, logger_filename):
+    global handler
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    # Set the format of root handlers
+    if not logging.getLogger().handlers:
+        logging.basicConfig(level=logging.INFO)
+    logging.getLogger().handlers[0].setFormatter(formatter)
+
+    # Redirect stdout and stderr to loggers
+    stdout_logger = logging.getLogger("stdout")
+    stdout_logger.setLevel(logging.INFO)
+    sl = StreamToLogger(stdout_logger, logging.INFO)
+    sys.stdout = sl
+
+    stderr_logger = logging.getLogger("stderr")
+    stderr_logger.setLevel(logging.ERROR)
+    sl = StreamToLogger(stderr_logger, logging.ERROR)
+    sys.stderr = sl
+
+    # Get logger
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    # Add a file handler for all loggers
+    if handler is None:
+        os.makedirs(LOGDIR, exist_ok=True)
+        filename = os.path.join(LOGDIR, logger_filename)
+        handler = logging.handlers.TimedRotatingFileHandler(
+            filename, when='D', utc=True, encoding='UTF-8')
+        handler.setFormatter(formatter)
+
+        for name, item in logging.root.manager.loggerDict.items():
+            if isinstance(item, logging.Logger):
+                item.addHandler(handler)
+
+    return logger
+
+
+class StreamToLogger(object):
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+    def __init__(self, logger, log_level=logging.INFO):
+        self.terminal = sys.stdout
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def __getattr__(self, attr):
+        return getattr(self.terminal, attr)
+
+    def write(self, buf):
+        temp_linebuf = self.linebuf + buf
+        self.linebuf = ''
+        for line in temp_linebuf.splitlines(True):
+            # From the io.TextIOWrapper docs:
+            #   On output, if newline is None, any '\n' characters written
+            #   are translated to the system default line separator.
+            # By default sys.stdout.write() expects '\n' newlines and then
+            # translates them so this is still cross platform.
+            if line[-1] == '\n':
+                self.logger.log(self.log_level, line.rstrip())
+            else:
+                self.linebuf += line
+
+    def flush(self):
+        if self.linebuf != '':
+            self.logger.log(self.log_level, self.linebuf.rstrip())
+        self.linebuf = ''
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def violates_moderation(text):
+    """
+    Check whether the text violates OpenAI moderation API.
+    """
+    url = "https://api.openai.com/v1/moderations"
+    headers = {"Content-Type": "application/json",
+               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}
+    text = text.replace("\n", "")
+    data = "{" + '"input": ' + f'"{text}"' + "}"
+    data = data.encode("utf-8")
+    try:
+        ret = requests.post(url, headers=headers, data=data, timeout=5)
+        flagged = ret.json()["results"][0]["flagged"]
+    except requests.exceptions.RequestException as e:
+        flagged = False
+    except KeyError as e:
+        flagged = False
+
+    return flagged
+
+
+def pretty_print_semaphore(semaphore):
+    if semaphore is None:
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"

requirements.txt

@@ -0,0 +1,23 @@
+# Core inference
+torch>=2.0.0
+torchvision>=0.15.0
+diffusers>=0.21.0
+transformers>=4.36.0
+accelerate>=0.25.0
+safetensors>=0.4.0
+
+# Image I/O
+pillow>=9.0.0
+tqdm>=4.65.0
+
+# LLaVA on-the-fly mode only (optional for reference-cache inference)
+peft>=0.7.0
+sentencepiece>=0.1.99
+protobuf>=3.20.0
+
+# LLaVA vision backbone
+timm>=0.9.0
+einops>=0.7.0
+
+# Hugging Face hub (weight download)
+huggingface-hub>=0.20.0

thera_llava.py

@@ -0,0 +1,21 @@
+"""
+Lazy loader for FrozenLLaVAExtractor.
+
+Reference-cache inference does not need LLaVA weights at runtime. The extractor
+is imported only when on-the-fly feature extraction is requested.
+"""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+
+def create_frozen_llava_extractor(**kwargs: Any):
+    """Construct FrozenLLaVAExtractor; imports LLaVA stack on first use."""
+    from llava.model.frozen_llava_extractor import FrozenLLaVAExtractor
+
+    sig = inspect.signature(FrozenLLaVAExtractor.__init__)
+    if "merge_lora" not in sig.parameters:
+        kwargs.pop("merge_lora", None)
+    return FrozenLLaVAExtractor(**kwargs)

thera_paths.py

@@ -0,0 +1,26 @@
+"""
+Default paths for TherA inference (relative to this repository root).
+
+Download model weights into the `weights/` directory — see README.md.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+WEIGHTS_DIR = PROJECT_ROOT / "weights"
+
+DEFAULT_CHECKPOINT = WEIGHTS_DIR / "checkpoint"
+DEFAULT_MERGED_MODEL = WEIGHTS_DIR / "merged_models"
+DEFAULT_PRETRAINED_SD = WEIGHTS_DIR / "stable-diffusion"
+DEFAULT_REFERENCE_CACHES = WEIGHTS_DIR / "reference_caches"
+
+
+def setup_project_path() -> Path:
+    """Ensure TherA root is on sys.path so local packages import correctly."""
+    root = str(PROJECT_ROOT)
+    if root not in sys.path:
+        sys.path.insert(0, root)
+    return PROJECT_ROOT

weights/.gitattributes

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

weights/README.md

@@ -0,0 +1,3 @@
+---
+license: mit
+---

weights/TherA_VLM/adapter_config.json

@@ -0,0 +1,34 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "./checkpoints/llava-1.5-7b-hf",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 256,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 128,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "up_proj",
+    "v_proj",
+    "q_proj",
+    "o_proj",
+    "gate_proj",
+    "down_proj",
+    "k_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

weights/TherA_VLM/config.json

@@ -0,0 +1,49 @@
+{
+  "_name_or_path": "./checkpoints/llava-1.5-7b-hf",
+  "architectures": [
+    "LlavaLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "freeze_mm_mlp_adapter": false,
+  "freeze_mm_vision_resampler": false,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "image_aspect_ratio": "pad",
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_length": 4096,
+  "max_position_embeddings": 4096,
+  "mm_hidden_size": 1024,
+  "mm_patch_merge_type": "flat",
+  "mm_projector_lr": 2e-05,
+  "mm_projector_type": "mlp2x_gelu",
+  "mm_resampler_type": null,
+  "mm_use_im_patch_token": false,
+  "mm_use_im_start_end": false,
+  "mm_vision_select_feature": "patch",
+  "mm_vision_select_layer": -2,
+  "mm_vision_tower": "openai/clip-vit-large-patch14-336",
+  "model_type": "llava_llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 32,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "tokenizer_model_max_length": 2048,
+  "tokenizer_padding_side": "right",
+  "torch_dtype": "float16",
+  "transformers_version": "4.37.2",
+  "tune_mm_mlp_adapter": false,
+  "tune_mm_vision_resampler": false,
+  "unfreeze_mm_vision_tower": false,
+  "use_cache": true,
+  "use_mm_proj": true,
+  "vocab_size": 32000
+}

weights/llava-1.5-7b-hf/.gitattributes

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

weights/llava-1.5-7b-hf/README.md

@@ -0,0 +1,155 @@
+---
+language:
+- en
+datasets:
+- liuhaotian/LLaVA-Instruct-150K
+pipeline_tag: image-text-to-text
+arxiv: 2304.08485
+license: llama2
+tags:
+- vision
+- image-text-to-text
+---
+# LLaVA Model Card
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/62441d1d9fdefb55a0b7d12c/FPshq08TKYD0e-qwPLDVO.png)
+
+Below is the model card of Llava model 7b, which is copied from the original Llava model card that you can find [here](https://huggingface.co/liuhaotian/llava-v1.5-13b).
+
+Check out also the Google Colab demo to run Llava on a free-tier Google Colab instance: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1qsl6cd2c8gGtEW1xV5io7S8NHh-Cp1TV?usp=sharing)
+
+Or check out our Spaces demo! [![Open in Spaces](https://huggingface.co/datasets/huggingface/badges/resolve/main/open-in-hf-spaces-md-dark.svg)](https://huggingface.co/spaces/llava-hf/llava-4bit)
+
+
+## Model details
+
+**Model type:**
+LLaVA is an open-source chatbot trained by fine-tuning LLaMA/Vicuna on GPT-generated multimodal instruction-following data.
+It is an auto-regressive language model, based on the transformer architecture.
+
+**Model date:**
+LLaVA-v1.5-7B was trained in September 2023.
+
+**Paper or resources for more information:**
+https://llava-vl.github.io/
+
+## How to use the model
+
+First, make sure to have `transformers >= 4.35.3`. 
+The model supports multi-image and multi-prompt generation. Meaning that you can pass multiple images in your prompt. Make sure also to follow the correct prompt template (`USER: xxx\nASSISTANT:`) and add the token `<image>` to the location where you want to query images:
+
+### Using `pipeline`:
+
+Below we used [`"llava-hf/llava-1.5-7b-hf"`](https://huggingface.co/llava-hf/llava-1.5-7b-hf) checkpoint.
+
+```python
+from transformers import pipeline
+
+pipe = pipeline("image-text-to-text", model="llava-hf/llava-1.5-7b-hf")
+messages = [
+    {
+      "role": "user",
+      "content": [
+          {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg"},
+          {"type": "text", "text": "What does the label 15 represent? (1) lava (2) core (3) tunnel (4) ash cloud"},
+        ],
+    },
+]
+
+out = pipe(text=messages, max_new_tokens=20)
+print(out)
+>>> [{'input_text': [{'role': 'user', 'content': [{'type': 'image', 'url': 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg'}, {'type': 'text', 'text': 'What does the label 15 represent? (1) lava (2) core (3) tunnel (4) ash cloud'}]}], 'generated_text': 'Lava'}]
+```
+
+### Using pure `transformers`:
+
+Below is an example script to run generation in `float16` precision on a GPU device:
+
+```python
+import requests
+from PIL import Image
+
+import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
+model_id = "llava-hf/llava-1.5-7b-hf"
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True, 
+).to(0)
+
+processor = AutoProcessor.from_pretrained(model_id)
+
+# Define a chat history and use `apply_chat_template` to get correctly formatted prompt
+# Each value in "content" has to be a list of dicts with types ("text", "image") 
+conversation = [
+    {
+
+      "role": "user",
+      "content": [
+          {"type": "text", "text": "What are these?"},
+          {"type": "image"},
+        ],
+    },
+]
+prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"
+raw_image = Image.open(requests.get(image_file, stream=True).raw)
+inputs = processor(images=raw_image, text=prompt, return_tensors='pt').to(0, torch.float16)
+
+output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
+print(processor.decode(output[0][2:], skip_special_tokens=True))
+```
+
+-----------
+From transformers>=v4.48, you can also pass image url or local path to the conversation history, and let the chat template handle the rest.
+Chat template will load the image for you and return inputs in `torch.Tensor` which you can pass directly to `model.generate()` 
+
+```python
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"}
+            {"type": "text", "text": "What is shown in this image?"},
+        ],
+    },
+]
+
+inputs = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors"pt")
+output = model.generate(**inputs, max_new_tokens=50)
+```
+
+### Model optimization
+
+#### 4-bit quantization through `bitsandbytes` library
+
+First make sure to install `bitsandbytes`, `pip install bitsandbytes` and make sure to have access to a CUDA compatible GPU device. Simply change the snippet above with: 
+
+```diff
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True,
++   load_in_4bit=True
+)
+```
+
+#### Use Flash-Attention 2 to further speed-up generation
+
+First make sure to install `flash-attn`. Refer to the [original repository of Flash Attention](https://github.com/Dao-AILab/flash-attention) regarding that package installation. Simply change the snippet above with: 
+
+```diff
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True,
++   use_flash_attention_2=True
+).to(0)
+```
+
+## License
+Llama 2 is licensed under the LLAMA 2 Community License, 
+Copyright (c) Meta Platforms, Inc. All Rights Reserved.

weights/llava-1.5-7b-hf/added_tokens.json

@@ -0,0 +1,4 @@
+{
+  "<image>": 32000,
+  "<pad>": 32001
+}

weights/llava-1.5-7b-hf/chat_template.jinja

@@ -0,0 +1,2 @@
+{% for message in messages %}{% if message['role'] != 'system' %}{{ message['role'].upper() + ': '}}{% endif %}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '<image>
+' }}{% endfor %}{# Render all text next #}{% if message['role'] != 'assistant' %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ content['text'] + ' '}}{% endfor %}{% else %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{% generation %}{{ content['text'] + ' '}}{% endgeneration %}{% endfor %}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:' }}{% endif %}

weights/llava-1.5-7b-hf/chat_template.json

@@ -0,0 +1,3 @@
+{
+  "chat_template": "{% for message in messages %}{% if message['role'] != 'system' %}{{ message['role'].upper() + ': '}}{% endif %}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '<image>\n' }}{% endfor %}{# Render all text next #}{% if message['role'] != 'assistant' %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ content['text'] + ' '}}{% endfor %}{% else %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{% generation %}{{ content['text'] + ' '}}{% endgeneration %}{% endfor %}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:' }}{% endif %}"
+}

weights/llava-1.5-7b-hf/config.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "LlavaForConditionalGeneration"
+  ],
+  "ignore_index": -100,
+  "image_token_index": 32000,
+  "model_type": "llava",
+  "pad_token_id": 32001,
+  "projector_hidden_act": "gelu",
+  "text_config": {
+    "_name_or_path": "lmsys/vicuna-7b-v1.5",
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "max_position_embeddings": 4096,
+    "model_type": "llama",
+    "rms_norm_eps": 1e-05,
+    "torch_dtype": "float16",
+    "vocab_size": 32064
+  },
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.36.0.dev0",
+  "vision_config": {
+    "hidden_size": 1024,
+    "image_size": 336,
+    "intermediate_size": 4096,
+    "model_type": "clip_vision_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "patch_size": 14,
+    "projection_dim": 768,
+    "vocab_size": 32000
+  },
+  "vision_feature_layer": -2,
+  "vision_feature_select_strategy": "default",
+  "vocab_size": 32064
+}

weights/llava-1.5-7b-hf/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 32001,
+  "transformers_version": "4.36.0.dev0"
+}

weights/llava-1.5-7b-hf/model.safetensors.index.json

@@ -0,0 +1,693 @@
+{
+  "metadata": {
+    "total_size": 14126854144
+  },
+  "weight_map": {
+    "language_model.lm_head.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.22.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.norm.weight": "model-00003-of-00003.safetensors",
+    "multi_modal_projector.linear_1.bias": "model-00001-of-00003.safetensors",
+    "multi_modal_projector.linear_1.weight": "model-00001-of-00003.safetensors",
+    "multi_modal_projector.linear_2.bias": "model-00001-of-00003.safetensors",
+    "multi_modal_projector.linear_2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.embeddings.class_embedding": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.embeddings.patch_embedding.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.embeddings.position_embedding.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.11.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.12.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.13.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.14.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.15.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.16.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.17.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.18.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.19.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.20.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.21.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.22.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.23.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.layer_norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.layer_norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.layer_norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.layer_norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.k_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.out_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.out_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.q_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.v_proj.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.encoder.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.post_layernorm.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.post_layernorm.weight": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.pre_layrnorm.bias": "model-00001-of-00003.safetensors",
+    "vision_tower.vision_model.pre_layrnorm.weight": "model-00001-of-00003.safetensors"
+  }
+}

weights/llava-1.5-7b-hf/preprocessor_config.json

@@ -0,0 +1,28 @@
+{
+  "crop_size": {
+    "height": 336,
+    "width": 336
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "processor_class": "LlavaProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 336
+  }
+}

weights/llava-1.5-7b-hf/processor_config.json

@@ -0,0 +1,7 @@
+{
+  "image_token": "<image>",
+  "num_additional_image_tokens": 1,
+  "patch_size": 14,
+  "processor_class": "LlavaProcessor",
+  "vision_feature_select_strategy": "default"
+}

weights/llava-1.5-7b-hf/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

weights/llava-1.5-7b-hf/tokenizer_config.json

@@ -0,0 +1,64 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {
+    "image_token": "<image>"
+  },
+  "image_token": "<image>",
+  "legacy": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "LlavaProcessor",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "trust_remote_code": false,
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

weights/merged_models/adapter/config.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "TextAdapter",
+  "_diffusers_version": "0.24.0",
+  "_name_or_path": "pretrained_weights/lavi-bridge/llama2_unet/adapter",
+  "in_dim": 4096,
+  "int_dim": 2432,
+  "out_dim": 768
+}

weights/merged_models/unet/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.24.0",
+  "_name_or_path": "pretrained_weights/stable-diffusion",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "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": 4,
+  "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": 4,
+  "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": 64,
+  "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": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

weights/stable-diffusion/scheduler/.ipynb_checkpoints/scheduler_config-checkpoint.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "PNDMScheduler",
+  "_diffusers_version": "0.2.2",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "num_train_timesteps": 1000,
+  "skip_prk_steps": true
+}

weights/stable-diffusion/scheduler/scheduler_config.json

@@ -0,0 +1,13 @@
+{
+  "_class_name": "PNDMScheduler",
+  "_diffusers_version": "0.7.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "num_train_timesteps": 1000,
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "trained_betas": null,
+  "clip_sample": false
+}

weights/stable-diffusion/vae/config.json

@@ -0,0 +1,29 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.2.2",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "out_channels": 3,
+  "sample_size": 512,
+  "scaling_factor": 0.18215,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}