[diffusion] model: support black-forest-labs/FLUX.2-dev (#14000)

python/sglang/multimodal_gen/configs/models/encoders/base.py

@@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 @dataclass
 class EncoderArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
     architectures: list[str] = field(default_factory=lambda: [])
     _supported_attention_backends: set[AttentionBackendEnum] = field(
         default_factory=lambda: {

python/sglang/multimodal_gen/configs/models/encoders/mistral.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m) -> bool:
+    return n.endswith("norm")
+
+
+@dataclass
+class Mistral3ArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = False
+    rope_theta: float = 10000.0
+    rope_scaling: float | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 256
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
+            (".gate_up_proj", ".up_proj", 1),  # type: ignore
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
+    )
+
+
+@dataclass
+class Mistral3Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=Mistral3ArchConfig)
+
+    prefix: str = "mistral3"

python/sglang/multimodal_gen/configs/models/vaes/flux.py

@@ -1,5 +1,3 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
@@ -28,6 +26,11 @@ class FluxVAEArchConfig(VAEArchConfig):
     clip_output: bool = True
 
 
+@dataclass
+class Flux2VAEArchConfig(FluxVAEArchConfig):
+    pass
+
+
 @dataclass
 class FluxVAEConfig(VAEConfig):
     arch_config: FluxVAEArchConfig = field(default_factory=FluxVAEArchConfig)
@@ -48,3 +51,8 @@ class FluxVAEConfig(VAEConfig):
             len(self.arch_config.block_out_channels) - 1
         )
         self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor
+
+
+@dataclass
+class Flux2VAEConfig(FluxVAEConfig):
+    arch_config: Flux2VAEArchConfig = field(default_factory=Flux2VAEArchConfig)

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -149,6 +149,8 @@ class PipelineConfig:
     preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
         default_factory=lambda: (preprocess_text,)
     )
+
+    # get prompt_embeds from encoder output
     postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], ...] = (
         field(default_factory=lambda: (postprocess_text,))
     )
@@ -175,15 +177,19 @@ class PipelineConfig:
     def slice_noise_pred(self, noise, latents):
         return noise
 
-    def adjust_size(self, width, height, image):
+    def maybe_resize_condition_image(self, width, height, image):
         """
         image: input image
         """
-        return width, height
+        return image, width, height
 
     def adjust_num_frames(self, num_frames):
         return num_frames
 
+    # tokenize the prompt
+    def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict:
+        return tokenizer(prompt, **tok_kwargs)
+
     # called in ImageEncodingStage, preprocess the image
     def preprocess_image(self, image, image_processor: VaeImageProcessor):
         return image
@@ -203,10 +209,32 @@ class PipelineConfig:
 
         return shape
 
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        vae_arch_config = self.vae_config.arch_config
+        scaling_factor = getattr(vae_arch_config, "scaling_factor", None)
+        if scaling_factor is None:
+            scaling_factor = getattr(vae, "scaling_factor", None)
+
+        shift_factor = getattr(vae_arch_config, "shift_factor", None)
+        if shift_factor is None:
+            shift_factor = getattr(vae, "shift_factor", None)
+        return scaling_factor, shift_factor
+
     # called after latents are prepared
     def maybe_pack_latents(self, latents, batch_size, batch):
         return latents
 
+    def maybe_prepare_latent_ids(self, latents):
+        return None
+
+    # called after vae encode
+    def post_process_vae_encode(self, image_latents, vae):
+        return image_latents
+
+    # called after scale_and_shift, before vae decoding
+    def preprocess_decoding(self, latents):
+        return latents
+
     def gather_latents_for_sp(self, latents):
         # For video latents [B, C, T_local, H, W], gather along time dim=2
         latents = sequence_model_parallel_all_gather(latents, dim=2)

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -1,8 +1,8 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-
+import math
 from dataclasses import dataclass, field
-from typing import Callable
+from typing import Callable, List, Optional
 
+import PIL.Image
 import torch
 
 from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
@@ -11,8 +11,15 @@ from sglang.multimodal_gen.configs.models.encoders import (
     BaseEncoderOutput,
     CLIPTextConfig,
     T5Config,
+    TextEncoderConfig,
 )
-from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
+from sglang.multimodal_gen.configs.models.encoders.base import TextEncoderArchConfig
+from sglang.multimodal_gen.configs.models.encoders.mistral import (
+    Mistral3Config,
+    _is_embeddings,
+    _is_transformer_layer,
+)
+from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig, FluxVAEConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
@@ -122,7 +129,7 @@ class FluxPipelineConfig(ImagePipelineConfig):
 
         return latent_image_ids
 
-    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb):
+    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
         txt_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device)
         img_ids = self._prepare_latent_image_ids(
             original_height=height,
@@ -156,7 +163,12 @@ class FluxPipelineConfig(ImagePipelineConfig):
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
         return {
             "freqs_cis": self.get_freqs_cis(
-                batch.prompt_embeds[1], batch.width, batch.height, device, rotary_emb
+                batch.prompt_embeds[1],
+                batch.width,
+                batch.height,
+                device,
+                rotary_emb,
+                batch,
             ),
             "pooled_projections": (
                 batch.pooled_embeds[0] if batch.pooled_embeds else None
@@ -171,8 +183,362 @@ class FluxPipelineConfig(ImagePipelineConfig):
                 batch.height,
                 device,
                 rotary_emb,
+                batch,
             ),
             "pooled_projections": (
                 batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None
             ),
         }
+
+
+def _prepare_latent_ids(
+    latents: torch.Tensor,  # (B, C, H, W)
+):
+    r"""
+    Generates 4D position coordinates (T, H, W, L) for latent tensors.
+
+    Args:
+        latents (torch.Tensor):
+            Latent tensor of shape (B, C, H, W)
+
+    Returns:
+        torch.Tensor:
+            Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0,
+            H=[0..H-1], W=[0..W-1], L=0
+    """
+
+    batch_size, _, height, width = latents.shape
+
+    t = torch.arange(1)  # [0] - time dimension
+    h = torch.arange(height)
+    w = torch.arange(width)
+    l = torch.arange(1)  # [0] - layer dimension
+
+    # Create position IDs: (H*W, 4)
+    latent_ids = torch.cartesian_prod(t, h, w, l)
+
+    # Expand to batch: (B, H*W, 4)
+    latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
+    return latent_ids
+
+
+def _unpack_latents_with_ids(
+    x: torch.Tensor, x_ids: torch.Tensor
+) -> list[torch.Tensor]:
+    """
+    using position ids to scatter tokens into place
+    """
+    x_list = []
+    x_ids = x_ids.to(device=x.device)
+    for data, pos in zip(x, x_ids):
+        _, ch = data.shape  # noqa: F841
+        h_ids = pos[:, 1].to(torch.int64)
+        w_ids = pos[:, 2].to(torch.int64)
+
+        h = torch.max(h_ids) + 1
+        w = torch.max(w_ids) + 1
+
+        flat_ids = h_ids * w + w_ids
+
+        out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype)
+        out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data)
+
+        # reshape from (H * W, C) to (H, W, C) and permute to (C, H, W)
+
+        out = out.view(h, w, ch).permute(2, 0, 1)
+        x_list.append(out)
+
+    return torch.stack(x_list, dim=0)
+
+
+def _patchify_latents(latents):
+    batch_size, num_channels_latents, height, width = latents.shape
+    latents = latents.view(
+        batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+    )
+    latents = latents.permute(0, 1, 3, 5, 2, 4)
+    latents = latents.reshape(
+        batch_size, num_channels_latents * 4, height // 2, width // 2
+    )
+    return latents
+
+
+def _unpatchify_latents(latents):
+    batch_size, num_channels_latents, height, width = latents.shape
+    latents = latents.reshape(
+        batch_size, num_channels_latents // (2 * 2), 2, 2, height, width
+    )
+    latents = latents.permute(0, 1, 4, 2, 5, 3)
+    latents = latents.reshape(
+        batch_size, num_channels_latents // (2 * 2), height * 2, width * 2
+    )
+    return latents
+
+
+def _prepare_text_ids(
+    x: torch.Tensor,  # (B, L, D) or (L, D)
+    t_coord: Optional[torch.Tensor] = None,
+):
+    B, L, _ = x.shape
+    out_ids = []
+
+    for i in range(B):
+        t = torch.arange(1) if t_coord is None else t_coord[i]
+        h = torch.arange(1)
+        w = torch.arange(1)
+        l = torch.arange(L)
+
+        coords = torch.cartesian_prod(t, h, w, l)
+        out_ids.append(coords)
+
+    return torch.stack(out_ids)
+
+
+def _prepare_image_ids(
+    image_latents: List[torch.Tensor],  # [(1, C, H, W), (1, C, H, W), ...]
+    scale: int = 10,
+):
+    if not isinstance(image_latents, list):
+        raise ValueError(
+            f"Expected `image_latents` to be a list, got {type(image_latents)}."
+        )
+
+    # create time offset for each reference image
+    t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))]
+    t_coords = [t.view(-1) for t in t_coords]
+
+    image_latent_ids = []
+    for x, t in zip(image_latents, t_coords):
+        x = x.squeeze(0)
+        _, height, width = x.shape
+
+        x_ids = torch.cartesian_prod(
+            t, torch.arange(height), torch.arange(width), torch.arange(1)
+        )
+        image_latent_ids.append(x_ids)
+
+    image_latent_ids = torch.cat(image_latent_ids, dim=0)
+    image_latent_ids = image_latent_ids.unsqueeze(0)
+
+    return image_latent_ids
+
+
+def flux_2_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
+    hidden_states_layers: list[int] = [10, 20, 30]
+
+    out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1)
+    batch_size, num_channels, seq_len, hidden_dim = out.shape
+    prompt_embeds = out.permute(0, 2, 1, 3).reshape(
+        batch_size, seq_len, num_channels * hidden_dim
+    )
+
+    return prompt_embeds
+
+
+@dataclass
+class Flux2MistralTextArchConfig(TextEncoderArchConfig):
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings]
+    )
+
+    def __post_init__(self):
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 512,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Flux2MistralTextConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=Flux2MistralTextArchConfig
+    )
+
+
+def format_text_input(prompts: List[str], system_message: str = None):
+    # Remove [IMG] tokens from prompts to avoid Pixtral validation issues
+    # when truncation is enabled. The processor counts [IMG] tokens and fails
+    # if the count changes after truncation.
+    cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts]
+
+    return [
+        [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": system_message}],
+            },
+            {"role": "user", "content": [{"type": "text", "text": prompt}]},
+        ]
+        for prompt in cleaned_txt
+    ]
+
+
+def flux_2_preprocess_text(prompt: str):
+    print(f"{prompt=}")
+    system_message = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation."
+    return format_text_input([prompt], system_message=system_message)
+
+
+# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
+def flux2_pack_latents(latents):
+    batch_size, num_channels, height, width = latents.shape
+    latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1)
+
+    return latents
+
+
+@dataclass
+class Flux2PipelineConfig(FluxPipelineConfig):
+    embedded_cfg_scale: float = 4.0
+
+    task_type: ModelTaskType = ModelTaskType.I2I
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Mistral3Config(),)
+    )
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Flux2MistralTextConfig(),)
+    )
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (flux_2_preprocess_text,),
+    )
+
+    postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (flux_2_postprocess_text,)
+    )
+    vae_config: VAEConfig = field(default_factory=Flux2VAEConfig)
+
+    def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
+        # flatten to 1-d list
+        prompts = [p for prompt in prompts for p in prompt]
+        inputs = tokenizer.apply_chat_template(
+            prompts,
+            add_generation_prompt=False,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            # 2048 from official github repo, 512 from diffusers
+            max_length=512,
+        )
+
+        return inputs
+
+    def prepare_latent_shape(self, batch, batch_size, num_frames):
+        height = 2 * (
+            batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
+        )
+        width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
+        num_channels_latents = self.dit_config.arch_config.in_channels
+        shape = (batch_size, num_channels_latents, height // 2, width // 2)
+        return shape
+
+    def get_pos_prompt_embeds(self, batch):
+        return batch.prompt_embeds[0]
+
+    def get_neg_prompt_embeds(self, batch):
+        return batch.negative_prompt_embeds[0]
+
+    def maybe_resize_condition_image(self, width, height, image):
+        target_area: int = 1024 * 1024
+
+        if width is not None and height is not None:
+            if width * height > target_area:
+                scale = math.sqrt(target_area / (width * height))
+                width = int(width * scale)
+                height = int(height * scale)
+                image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+                width, height = image.size
+
+        return image, width, height
+
+    def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
+
+        txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)
+
+        img_ids = batch.latent_ids
+        if batch.image_latent is not None:
+            image_latent_ids = batch.condition_image_latent_ids
+            img_ids = torch.cat([img_ids, image_latent_ids], dim=1).to(device=device)
+
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        # NOTE(mick): prepare it here, to avoid unnecessary computations
+        img_cos, img_sin = rotary_emb.forward(img_ids)
+        img_cos = shard_rotary_emb_for_sp(img_cos)
+        img_sin = shard_rotary_emb_for_sp(img_sin)
+
+        txt_cos, txt_sin = rotary_emb.forward(txt_ids)
+
+        cos = torch.cat([txt_cos, img_cos], dim=0).to(device=device)
+        sin = torch.cat([txt_sin, img_sin], dim=0).to(device=device)
+        return cos, sin
+
+    def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        return {
+            "freqs_cis": self.get_freqs_cis(
+                batch.prompt_embeds[0],
+                batch.width,
+                batch.height,
+                device,
+                rotary_emb,
+                batch,
+            )
+        }
+
+    def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        return {}
+
+    def maybe_pack_latents(self, latents, batch_size, batch):
+        return flux2_pack_latents(latents)
+
+    def maybe_prepare_latent_ids(self, latents):
+        return _prepare_latent_ids(latents)
+
+    def post_process_vae_encode(self, image_latents, vae):
+        # patchify
+        image_latents = _patchify_latents(image_latents)
+        return image_latents
+
+    def preprocess_decoding(self, latents):
+        latents = _unpatchify_latents(latents)
+        return latents
+
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        vae_arch_config = self.vae_config.arch_config
+        latents_bn_mean = (
+            vae.bn.running_mean.view(1, -1, 1, 1).to(device=device).to(device, dtype)
+        )
+        latents_bn_std = torch.sqrt(
+            vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps
+        ).to(device, dtype)
+        return 1 / latents_bn_std, latents_bn_mean
+
+    def post_denoising_loop(self, latents, batch):
+        latent_ids = batch.latent_ids
+        latents = _unpack_latents_with_ids(latents, latent_ids)
+
+        return latents
+
+    def slice_noise_pred(self, noise, latents):
+        # remove noise over input image
+        noise = noise[:, : latents.size(1) :]
+        return noise

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -123,6 +123,18 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         # pack latents
         return _pack_latents(latents, batch_size, num_channels_latents, height, width)
 
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        vae_arch_config = self.vae_config.arch_config
+        scaling_factor = 1.0 / torch.tensor(
+            vae_arch_config.latents_std, device=device
+        ).view(1, vae_arch_config.z_dim, 1, 1, 1).to(device, dtype)
+        shift_factor = (
+            torch.tensor(vae_arch_config.latents_mean)
+            .view(1, vae_arch_config.z_dim, 1, 1, 1)
+            .to(device, dtype)
+        )
+        return scaling_factor, shift_factor
+
     @staticmethod
     def get_freqs_cis(img_shapes, txt_seq_lens, rotary_emb, device, dtype):
         # img_shapes: for global entire image
@@ -203,7 +215,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         assert batch_size == 1
         height = batch.height
         width = batch.width
-        image = batch.pil_image
+        image = batch.condition_image
         image_size = image[0].size if isinstance(image, list) else image.size
         edit_width, edit_height, _ = calculate_dimensions(
             1024 * 1024, image_size[0] / image_size[1]
@@ -268,7 +280,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         image = image_processor.resize(image, calculated_height, calculated_width)
         return image
 
-    def adjust_size(self, width, height, image):
+    def maybe_resize_condition_image(self, width, height, image):
         image_size = image[0].size if isinstance(image, list) else image.size
         calculated_width, calculated_height, _ = calculate_dimensions(
             1024 * 1024, image_size[0] / image_size[1]

python/sglang/multimodal_gen/configs/sample/base.py

@@ -315,8 +315,10 @@ class SamplingParams:
         sampling_params = SamplingParams.from_pretrained(model_path)
 
         user_sampling_params = SamplingParams(*args, **kwargs)
+        # TODO: refactor
         sampling_params._merge_with_user_params(user_sampling_params)
-
+        sampling_params.width_not_provided = user_sampling_params.width is None
+        sampling_params.height_not_provided = user_sampling_params.height is None
         sampling_params.adjust(server_args)
 
         return sampling_params
@@ -539,6 +541,8 @@ class SamplingParams:
                 if hasattr(self, field_name):
                     setattr(self, field_name, user_value)
 
+        self.height_not_provided = user_params.height_not_provided
+        self.width_not_provided = user_params.width_not_provided
         self.__post_init__()
 
     @property

python/sglang/multimodal_gen/configs/sample/flux.py

@@ -16,3 +16,10 @@ class FluxSamplingParams(SamplingParams):
     guidance_scale: float = 1.0
     negative_prompt: str = None
     num_inference_steps: int = 50
+
+    def __post_init__(self):
+        default_sample_size = 128
+        vae_scale_factor = 8
+        # FIXME
+        # self.height = default_sample_size * vae_scale_factor
+        # self.width = default_sample_size * vae_scale_factor