Add img2img, more options, gradio interface

.gitignore

@@ -1 +1,12 @@
 __pycache__
+*.jpg
+*.png
+*.jpeg
+*.gif
+*.bmp
+*.webp
+*.mp4
+*.mp3
+*.mp3
+*.txt
+.copilotignore

api.py

@@ -17,6 +17,8 @@ class GenerateArgs(BaseModel):
     seed: Optional[int] = Field(
         default_factory=lambda: np.random.randint(0, 2**32 - 1), gt=0, lt=2**32 - 1
     )
+    strength: Optional[float] = 1.0
+    init_image: Optional[str] = None
 
 
 @app.post("/generate")

configs/config-dev-gigaquant.json

@@ -0,0 +1,52 @@
+{
+  "version": "flux-dev",
+  "params": {
+    "in_channels": 64,
+    "vec_in_dim": 768,
+    "context_in_dim": 4096,
+    "hidden_size": 3072,
+    "mlp_ratio": 4.0,
+    "num_heads": 24,
+    "depth": 19,
+    "depth_single_blocks": 38,
+    "axes_dim": [
+      16,
+      56,
+      56
+    ],
+    "theta": 10000,
+    "qkv_bias": true,
+    "guidance_embed": true
+  },
+  "ae_params": {
+    "resolution": 256,
+    "in_channels": 3,
+    "ch": 128,
+    "out_ch": 3,
+    "ch_mult": [
+      1,
+      2,
+      4,
+      4
+    ],
+    "num_res_blocks": 2,
+    "z_channels": 16,
+    "scale_factor": 0.3611,
+    "shift_factor": 0.1159
+  },
+  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
+  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
+  "repo_id": "black-forest-labs/FLUX.1-dev",
+  "repo_flow": "flux1-dev.sft",
+  "repo_ae": "ae.sft",
+  "text_enc_max_length": 512,
+  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
+  "text_enc_device": "cuda:1",
+  "ae_device": "cuda:1",
+  "flux_device": "cuda:0",
+  "flow_dtype": "float16",
+  "ae_dtype": "bfloat16",
+  "text_enc_dtype": "bfloat16",
+  "num_to_quant": 8000,
+  "quantize_extras": true
+}

configs/config-dev.json

@@ -47,5 +47,7 @@
   "flow_dtype": "float16",
   "ae_dtype": "bfloat16",
   "text_enc_dtype": "bfloat16",
-  "num_to_quant": 20
+  "num_to_quant": 22,
+  "compile_extras": false,
+  "compile_blocks": false
 }

cublas_linear.py

@@ -1,152 +1 @@
-import math
-from typing import Literal, Optional
-
-import torch
-from torch.nn import functional as F
-
-from cublas_ops_ext import _simt_hgemv
-from cublas_ops_ext import cublas_hgemm_axbT as _cublas_hgemm_axbT
-from cublas_ops_ext import cublas_hgemm_batched_simple as _cublas_hgemm_batched_simple
-from cublas_ops_ext import (
-    cublaslt_hgemm_batched_simple as _cublaslt_hgemm_batched_simple,
-)
-from cublas_ops_ext import cublaslt_hgemm_simple as _cublaslt_hgemm_simple
-from torch import Tensor, nn
-
-global has_moved
-has_moved = {idx: False for idx in range(torch.cuda.device_count())}
-
-
-class StaticState:
-    workspace = {
-        idx: torch.empty((1024 * 1024 * 8,), dtype=torch.uint8)
-        for idx in range(torch.cuda.device_count())
-    }
-    workspace_size = workspace[0].nelement()
-    bias_g = {
-        idx: torch.tensor([], dtype=torch.float16)
-        for idx in range(torch.cuda.device_count())
-    }
-
-    @classmethod
-    def get(cls, __name: str, device: torch.device) -> torch.Any:
-        global has_moved
-        idx = device.index if device.index is not None else 0
-        if not has_moved[idx]:
-            cls.workspace[idx] = cls.workspace[idx].cuda(idx)
-            cls.bias_g[idx] = cls.bias_g[idx].cuda(idx)
-            has_moved[idx] = True
-        if "bias" in __name:
-            return cls.bias_g[idx]
-        if "workspace" in __name:
-            return cls.workspace[idx]
-        if "workspace_size" in __name:
-            return cls.workspace_size
-
-
-@torch.no_grad()
-def hgemv_simt(vec: torch.HalfTensor, mat: torch.HalfTensor, block_dim_x: int = 32):
-    prev_dims = vec.shape[:-1]
-    out = _simt_hgemv(mat, vec.view(-1, 1), block_dim_x=block_dim_x).view(
-        *prev_dims, -1
-    )
-    return out
-
-
-@torch.no_grad()
-def cublas_half_matmul_batched_simple(a: torch.Tensor, b: torch.Tensor):
-    out = _cublas_hgemm_batched_simple(a, b)
-    return out
-
-
-@torch.no_grad()
-def cublas_half_matmul_simple(a: torch.Tensor, b: torch.Tensor):
-    out = _cublas_hgemm_axbT(b, a)
-    return out
-
-
-@torch.no_grad()
-def cublaslt_fused_half_matmul_simple(
-    a: torch.Tensor,
-    b: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
-    epilogue_str: Optional[Literal["NONE", "RELU", "GELU"]] = "NONE",
-):
-    if bias is None:
-        bias = StaticState.get("bias", a.device)
-    out = _cublaslt_hgemm_simple(
-        a, b, bias, epilogue_str, StaticState.get("workspace", a.device)
-    )
-    return out
-
-
-@torch.no_grad()
-def cublaslt_fused_half_matmul_batched_simple(
-    a: torch.Tensor,
-    b: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
-    epilogue_str: Optional[Literal["NONE", "RELU", "GELU"]] = "NONE",
-):
-    if bias is None:
-        bias = StaticState.get("bias", a.device)
-    out = _cublaslt_hgemm_batched_simple(
-        a, b, bias, epilogue_str, StaticState.get("workspace", a.device)
-    )
-    return out
-
-
-class CublasLinear(nn.Linear):
-    def __init__(
-        self,
-        in_features,
-        out_features,
-        bias=True,
-        device=None,
-        dtype=torch.float16,
-        epilogue_str="NONE",
-    ):
-        super().__init__(
-            in_features, out_features, bias=bias, device=device, dtype=dtype
-        )
-        self._epilogue_str = epilogue_str
-        self.has_bias = bias
-        self.has_checked_weight = False
-
-    def forward(self, x: Tensor) -> Tensor:
-        if not self.has_checked_weight:
-            if not self.weight.dtype == torch.float16:
-                self.to(dtype=torch.float16)
-            self.has_checked_weight = True
-        out_dtype = x.dtype
-        needs_convert = out_dtype != torch.float16
-        if needs_convert:
-            x = x.type(torch.float16)
-
-        use_cublasLt = self.has_bias or self._epilogue_str != "NONE"
-        if x.ndim == 1:
-            x = x.unsqueeze(0)
-        if math.prod(x.shape) == x.shape[-1]:
-            out = F.linear(x, self.weight, bias=self.bias)
-            if self._epilogue_str == "RELU":
-                return F.relu(out)
-            elif self._epilogue_str == "GELU":
-                return F.gelu(out)
-            if needs_convert:
-                return out.type(out_dtype)
-            return out
-        if use_cublasLt:
-            leading_dims = x.shape[:-1]
-            x = x.reshape(-1, x.shape[-1])
-            out = cublaslt_fused_half_matmul_simple(
-                x, self.weight, bias=self.bias.data, epilogue_str=self._epilogue_str
-            )
-            if needs_convert:
-                return out.view(*leading_dims, out.shape[-1]).type(out_dtype)
-            return out.view(*leading_dims, out.shape[-1])
-        else:
-            leading_dims = x.shape[:-1]
-            x = x.reshape(-1, x.shape[-1])
-            out = cublas_half_matmul_simple(x, self.weight)
-            if needs_convert:
-                return out.view(*leading_dims, out.shape[-1]).type(out_dtype)
-            return out.view(*leading_dims, out.shape[-1])
+from cublas_ops import CublasLinear

flux_emphasis.py

@@ -0,0 +1,448 @@
+from typing import TYPE_CHECKING, Optional
+from pydash import flatten
+
+import torch
+from transformers.models.clip.tokenization_clip import CLIPTokenizer
+from einops import repeat
+
+if TYPE_CHECKING:
+    from flux_pipeline import FluxPipeline
+
+
+def parse_prompt_attention(text):
+    """
+    Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
+    Accepted tokens are:
+      (abc) - increases attention to abc by a multiplier of 1.1
+      (abc:3.12) - increases attention to abc by a multiplier of 3.12
+      [abc] - decreases attention to abc by a multiplier of 1.1
+      \\( - literal character '('
+      \\[ - literal character '['
+      \\) - literal character ')'
+      \\] - literal character ']'
+      \\ - literal character '\'
+      anything else - just text
+
+    >>> parse_prompt_attention('normal text')
+    [['normal text', 1.0]]
+    >>> parse_prompt_attention('an (important) word')
+    [['an ', 1.0], ['important', 1.1], [' word', 1.0]]
+    >>> parse_prompt_attention('(unbalanced')
+    [['unbalanced', 1.1]]
+    >>> parse_prompt_attention('\\(literal\\]')
+    [['(literal]', 1.0]]
+    >>> parse_prompt_attention('(unnecessary)(parens)')
+    [['unnecessaryparens', 1.1]]
+    >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
+    [['a ', 1.0],
+     ['house', 1.5730000000000004],
+     [' ', 1.1],
+     ['on', 1.0],
+     [' a ', 1.1],
+     ['hill', 0.55],
+     [', sun, ', 1.1],
+     ['sky', 1.4641000000000006],
+     ['.', 1.1]]
+    """
+    import re
+
+    re_attention = re.compile(
+        r"""
+            \\\(|\\\)|\\\[|\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|
+            \)|]|[^\\()\[\]:]+|:
+        """,
+        re.X,
+    )
+
+    re_break = re.compile(r"\s*\bBREAK\b\s*", re.S)
+
+    res = []
+    round_brackets = []
+    square_brackets = []
+
+    round_bracket_multiplier = 1.1
+    square_bracket_multiplier = 1 / 1.1
+
+    def multiply_range(start_position, multiplier):
+        for p in range(start_position, len(res)):
+            res[p][1] *= multiplier
+
+    for m in re_attention.finditer(text):
+        text = m.group(0)
+        weight = m.group(1)
+
+        if text.startswith("\\"):
+            res.append([text[1:], 1.0])
+        elif text == "(":
+            round_brackets.append(len(res))
+        elif text == "[":
+            square_brackets.append(len(res))
+        elif weight is not None and len(round_brackets) > 0:
+            multiply_range(round_brackets.pop(), float(weight))
+        elif text == ")" and len(round_brackets) > 0:
+            multiply_range(round_brackets.pop(), round_bracket_multiplier)
+        elif text == "]" and len(square_brackets) > 0:
+            multiply_range(square_brackets.pop(), square_bracket_multiplier)
+        else:
+            parts = re.split(re_break, text)
+            for i, part in enumerate(parts):
+                if i > 0:
+                    res.append(["BREAK", -1])
+                res.append([part, 1.0])
+
+    for pos in round_brackets:
+        multiply_range(pos, round_bracket_multiplier)
+
+    for pos in square_brackets:
+        multiply_range(pos, square_bracket_multiplier)
+
+    if len(res) == 0:
+        res = [["", 1.0]]
+
+    # merge runs of identical weights
+    i = 0
+    while i + 1 < len(res):
+        if res[i][1] == res[i + 1][1]:
+            res[i][0] += res[i + 1][0]
+            res.pop(i + 1)
+        else:
+            i += 1
+
+    return res
+
+
+def get_prompts_tokens_with_weights(clip_tokenizer: CLIPTokenizer, prompt: str):
+    """
+    Get prompt token ids and weights, this function works for both prompt and negative prompt
+
+    Args:
+        pipe (CLIPTokenizer)
+            A CLIPTokenizer
+        prompt (str)
+            A prompt string with weights
+
+    Returns:
+        text_tokens (list)
+            A list contains token ids
+        text_weight (list)
+            A list contains the correspodent weight of token ids
+
+    Example:
+        import torch
+        from transformers import CLIPTokenizer
+
+        clip_tokenizer = CLIPTokenizer.from_pretrained(
+            "stablediffusionapi/deliberate-v2"
+            , subfolder = "tokenizer"
+            , dtype = torch.float16
+        )
+
+        token_id_list, token_weight_list = get_prompts_tokens_with_weights(
+            clip_tokenizer = clip_tokenizer
+            ,prompt = "a (red:1.5) cat"*70
+        )
+    """
+    texts_and_weights = parse_prompt_attention(prompt)
+    text_tokens, text_weights = [], []
+    maxlen = clip_tokenizer.model_max_length
+    for word, weight in texts_and_weights:
+        # tokenize and discard the starting and the ending token
+        token = clip_tokenizer(
+            word, truncation=False, padding=False, add_special_tokens=False
+        ).input_ids
+        # so that tokenize whatever length prompt
+        # the returned token is a 1d list: [320, 1125, 539, 320]
+        print(
+            token,
+            "|FOR MODEL LEN{}|".format(maxlen),
+            clip_tokenizer.decode(
+                token, skip_special_tokens=True, clean_up_tokenization_spaces=True
+            ),
+        )
+        # merge the new tokens to the all tokens holder: text_tokens
+        text_tokens = [*text_tokens, *token]
+
+        # each token chunk will come with one weight, like ['red cat', 2.0]
+        # need to expand weight for each token.
+        chunk_weights = [weight] * len(token)
+
+        # append the weight back to the weight holder: text_weights
+        text_weights = [*text_weights, *chunk_weights]
+    return text_tokens, text_weights
+
+
+def group_tokens_and_weights(
+    token_ids: list,
+    weights: list,
+    pad_last_block=False,
+    bos=49406,
+    eos=49407,
+    max_length=77,
+    pad_tokens=True,
+):
+    """
+    Produce tokens and weights in groups and pad the missing tokens
+
+    Args:
+        token_ids (list)
+            The token ids from tokenizer
+        weights (list)
+            The weights list from function get_prompts_tokens_with_weights
+        pad_last_block (bool)
+            Control if fill the last token list to 75 tokens with eos
+    Returns:
+        new_token_ids (2d list)
+        new_weights (2d list)
+
+    Example:
+        token_groups,weight_groups = group_tokens_and_weights(
+            token_ids = token_id_list
+            , weights = token_weight_list
+        )
+    """
+    max_len = max_length - 2 if max_length < 77 else max_length
+    # this will be a 2d list
+    new_token_ids = []
+    new_weights = []
+    while len(token_ids) >= max_len:
+        # get the first 75 tokens
+        head_75_tokens = [token_ids.pop(0) for _ in range(max_len)]
+        head_75_weights = [weights.pop(0) for _ in range(max_len)]
+
+        # extract token ids and weights
+
+        if pad_tokens:
+            if bos is not None:
+                temp_77_token_ids = [bos] + head_75_tokens + [eos]
+                temp_77_weights = [1.0] + head_75_weights + [1.0]
+            else:
+                temp_77_token_ids = head_75_tokens + [eos]
+                temp_77_weights = head_75_weights + [1.0]
+
+        # add 77 token and weights chunk to the holder list
+        new_token_ids.append(temp_77_token_ids)
+        new_weights.append(temp_77_weights)
+
+    # padding the left
+    if len(token_ids) > 0:
+        if pad_tokens:
+            padding_len = max_len - len(token_ids) if pad_last_block else 0
+
+            temp_77_token_ids = [bos] + token_ids + [eos] * padding_len + [eos]
+            new_token_ids.append(temp_77_token_ids)
+
+            temp_77_weights = [1.0] + weights + [1.0] * padding_len + [1.0]
+            new_weights.append(temp_77_weights)
+        else:
+            new_token_ids.append(token_ids)
+            new_weights.append(weights)
+    return new_token_ids, new_weights
+
+
+def standardize_tensor(
+    input_tensor: torch.Tensor, target_mean: float, target_std: float
+) -> torch.Tensor:
+    """
+    This function standardizes an input tensor so that it has a specific mean and standard deviation.
+
+    Parameters:
+    input_tensor (torch.Tensor): The tensor to standardize.
+    target_mean (float): The target mean for the tensor.
+    target_std (float): The target standard deviation for the tensor.
+
+    Returns:
+    torch.Tensor: The standardized tensor.
+    """
+
+    # First, compute the mean and std of the input tensor
+    mean = input_tensor.mean()
+    std = input_tensor.std()
+
+    # Then, standardize the tensor to have a mean of 0 and std of 1
+    standardized_tensor = (input_tensor - mean) / std
+
+    # Finally, scale the tensor to the target mean and std
+    output_tensor = standardized_tensor * target_std + target_mean
+
+    return output_tensor
+
+
+def apply_weights(
+    prompt_tokens: torch.Tensor,
+    weight_tensor: torch.Tensor,
+    token_embedding: torch.Tensor,
+    eos_token_id: int,
+    pad_last_block: bool = True,
+) -> torch.FloatTensor:
+    mean = token_embedding.mean()
+    std = token_embedding.std()
+    if pad_last_block:
+        pooled_tensor = token_embedding[
+            torch.arange(token_embedding.shape[0], device=token_embedding.device),
+            (
+                prompt_tokens.to(dtype=torch.int, device=token_embedding.device)
+                == eos_token_id
+            )
+            .int()
+            .argmax(dim=-1),
+        ]
+    else:
+        pooled_tensor = token_embedding[:, -1]
+
+    for j in range(len(weight_tensor)):
+        if weight_tensor[j] != 1.0:
+            token_embedding[:, j] = (
+                pooled_tensor
+                + (token_embedding[:, j] - pooled_tensor) * weight_tensor[j]
+            )
+    return standardize_tensor(token_embedding, mean, std)
+
+
+@torch.inference_mode()
+def get_weighted_text_embeddings_flux(
+    pipe: "FluxPipeline",
+    prompt: str = "",
+    num_images_per_prompt: int = 1,
+    device: Optional[torch.device] = None,
+    target_device: Optional[torch.device] = torch.device("cuda:0"),
+    target_dtype: Optional[torch.dtype] = torch.bfloat16,
+):
+    """
+    This function can process long prompt with weights, no length limitation
+    for Stable Diffusion XL
+
+    Args:
+        pipe (StableDiffusionPipeline)
+        prompt (str)
+        prompt_2 (str)
+        neg_prompt (str)
+        neg_prompt_2 (str)
+        num_images_per_prompt (int)
+        device (torch.device)
+    Returns:
+        prompt_embeds (torch.Tensor)
+        neg_prompt_embeds (torch.Tensor)
+    """
+    device = device or pipe._execution_device
+
+    eos = pipe.clip.tokenizer.eos_token_id
+    eos_2 = pipe.t5.tokenizer.eos_token_id
+    bos = pipe.clip.tokenizer.bos_token_id
+    bos_2 = pipe.t5.tokenizer.bos_token_id
+
+    clip = pipe.clip.hf_module
+    t5 = pipe.t5.hf_module
+
+    tokenizer_clip = pipe.clip.tokenizer
+    tokenizer_t5 = pipe.t5.tokenizer
+
+    t5_length = 512 if pipe.name == "flux-dev" else 256
+    clip_length = 77
+
+    tokenizer_t5(
+        prompt,
+        add_special_tokens=True,
+        padding="max_length",
+        truncation=True,
+        max_length=t5_length,
+        return_tensors="pt",
+    )
+
+    # tokenizer 1
+    prompt_tokens_clip, prompt_weights_clip = get_prompts_tokens_with_weights(
+        tokenizer_clip, prompt
+    )
+
+    # tokenizer 2
+    prompt_tokens_t5, prompt_weights_t5 = get_prompts_tokens_with_weights(
+        tokenizer_t5, prompt
+    )
+
+    prompt_tokens_clip_grouped, prompt_weights_clip_grouped = group_tokens_and_weights(
+        prompt_tokens_clip,
+        prompt_weights_clip,
+        pad_last_block=True,
+        bos=bos,
+        eos=eos,
+        max_length=clip_length,
+    )
+    prompt_tokens_t5_grouped, prompt_weights_t5_grouped = group_tokens_and_weights(
+        prompt_tokens_t5,
+        prompt_weights_t5,
+        pad_last_block=True,
+        bos=bos_2,
+        eos=eos_2,
+        max_length=t5_length,
+        pad_tokens=False,
+    )
+    prompt_tokens_t5 = flatten(prompt_tokens_t5_grouped)
+    prompt_weights_t5 = flatten(prompt_weights_t5_grouped)
+    prompt_tokens_clip = flatten(prompt_tokens_clip_grouped)
+    prompt_weights_clip = flatten(prompt_weights_clip_grouped)
+
+    prompt_tokens_clip = tokenizer_clip.decode(
+        prompt_tokens_clip, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )
+    prompt_tokens_clip = tokenizer_clip(
+        prompt_tokens_clip,
+        add_special_tokens=True,
+        padding="max_length",
+        truncation=True,
+        max_length=clip_length,
+        return_tensors="pt",
+    ).input_ids.to(device)
+    prompt_tokens_t5 = tokenizer_t5.decode(
+        prompt_tokens_t5, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )
+    prompt_tokens_t5 = tokenizer_t5(
+        prompt_tokens_t5,
+        add_special_tokens=True,
+        padding="max_length",
+        truncation=True,
+        max_length=t5_length,
+        return_tensors="pt",
+    ).input_ids.to(device)
+
+    prompt_weights_t5 = torch.cat(
+        [
+            torch.tensor(prompt_weights_t5, dtype=torch.float32),
+            torch.full(
+                (t5_length - torch.tensor(prompt_weights_t5).numel(),),
+                1.0,
+                dtype=torch.float32,
+            ),
+        ],
+        dim=0,
+    ).to(device)
+
+    clip_embeds = clip(
+        prompt_tokens_clip, output_hidden_states=True, attention_mask=None
+    )["pooler_output"]
+    if clip_embeds.shape[0] == 1 and num_images_per_prompt > 1:
+        clip_embeds = repeat(clip_embeds, "1 ... -> bs ...", bs=num_images_per_prompt)
+
+    weight_tensor_t5 = torch.tensor(
+        flatten(prompt_weights_t5), dtype=torch.float32, device=device
+    )
+    t5_embeds = t5(prompt_tokens_t5, output_hidden_states=True, attention_mask=None)[
+        "last_hidden_state"
+    ]
+    t5_embeds = apply_weights(prompt_tokens_t5, weight_tensor_t5, t5_embeds, eos_2)
+    print(t5_embeds.shape)
+    if t5_embeds.shape[0] == 1 and num_images_per_prompt > 1:
+        t5_embeds = repeat(t5_embeds, "1 ... -> bs ...", bs=num_images_per_prompt)
+    txt_ids = torch.zeros(
+        num_images_per_prompt,
+        t5_embeds.shape[1],
+        3,
+        device=target_device,
+        dtype=target_dtype,
+    )
+    t5_embeds = t5_embeds.to(target_device, dtype=target_dtype)
+    clip_embeds = clip_embeds.to(target_device, dtype=target_dtype)
+
+    return (
+        clip_embeds,
+        t5_embeds,
+        txt_ids,
+    )

flux_impl.py

@@ -1,272 +0,0 @@
-import io
-from typing import List
-
-import torch
-from torch import nn
-
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
-torch.backends.cudnn.benchmark_limit = 20
-torch.set_float32_matmul_precision("high")
-from torch._dynamo import config
-from torch._inductor import config as ind_config
-
-config.cache_size_limit = 10000000000
-ind_config.force_fuse_int_mm_with_mul = True
-
-from loguru import logger
-from torchao.quantization.quant_api import int8_weight_only, quantize_
-
-from cublas_linear import CublasLinear as F16Linear
-from modules.flux_model import RMSNorm
-from sampling import denoise, get_noise, get_schedule, prepare, unpack
-from turbojpeg_imgs import TurboImage
-from util import (
-    ModelSpec,
-    into_device,
-    into_dtype,
-    load_config_from_path,
-    load_models_from_config,
-)
-
-
-class Model:
-    def __init__(
-        self,
-        name,
-        offload=False,
-        clip=None,
-        t5=None,
-        model=None,
-        ae=None,
-        dtype=torch.bfloat16,
-        verbose=False,
-        flux_device="cuda:0",
-        ae_device="cuda:1",
-        clip_device="cuda:1",
-        t5_device="cuda:1",
-    ):
-
-        self.name = name
-        self.device_flux = (
-            flux_device
-            if isinstance(flux_device, torch.device)
-            else torch.device(flux_device)
-        )
-        self.device_ae = (
-            ae_device
-            if isinstance(ae_device, torch.device)
-            else torch.device(ae_device)
-        )
-        self.device_clip = (
-            clip_device
-            if isinstance(clip_device, torch.device)
-            else torch.device(clip_device)
-        )
-        self.device_t5 = (
-            t5_device
-            if isinstance(t5_device, torch.device)
-            else torch.device(t5_device)
-        )
-        self.dtype = dtype
-        self.offload = offload
-        self.clip = clip
-        self.t5 = t5
-        self.model = model
-        self.ae = ae
-        self.rng = torch.Generator(device="cpu")
-        self.turbojpeg = TurboImage()
-        self.verbose = verbose
-
-    @torch.inference_mode()
-    def generate(
-        self,
-        prompt,
-        width=720,
-        height=1023,
-        num_steps=24,
-        guidance=3.5,
-        seed=None,
-    ):
-        if num_steps is None:
-            num_steps = 4 if self.name == "flux-schnell" else 50
-
-        # allow for packing and conversion to latent space
-        height = 16 * (height // 16)
-        width = 16 * (width // 16)
-
-        if seed is None:
-            seed = self.rng.seed()
-        logger.info(f"Generating with:\nSeed: {seed}\nPrompt: {prompt}")
-
-        x = get_noise(
-            1,
-            height,
-            width,
-            device=self.device_t5,
-            dtype=torch.bfloat16,
-            seed=seed,
-        )
-        inp = prepare(self.t5, self.clip, x, prompt=prompt)
-        timesteps = get_schedule(
-            num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell")
-        )
-        for k in inp:
-            inp[k] = inp[k].to(self.device_flux).type(self.dtype)
-
-        # denoise initial noise
-        x = denoise(
-            self.model,
-            **inp,
-            timesteps=timesteps,
-            guidance=guidance,
-            dtype=self.dtype,
-            device=self.device_flux,
-        )
-        inp.clear()
-        timesteps.clear()
-        torch.cuda.empty_cache()
-        x = x.to(self.device_ae)
-
-        # decode latents to pixel space
-        x = unpack(x.float(), height, width)
-        with torch.autocast(
-            device_type=self.device_ae.type, dtype=torch.bfloat16, cache_enabled=False
-        ):
-            x = self.ae.decode(x)
-
-        # bring into PIL format and save
-        x = x.clamp(-1, 1)
-        num_images = x.shape[0]
-        images: List[torch.Tensor] = []
-        for i in range(num_images):
-            x = x[i].permute(1, 2, 0).add(1.0).mul(127.5).type(torch.uint8).contiguous()
-            images.append(x)
-        if len(images) == 1:
-            im = images[0]
-        else:
-            im = torch.vstack(images)
-
-        im = self.turbojpeg.encode_torch(im, quality=95)
-        images.clear()
-        return io.BytesIO(im)
-
-
-def quant_module(module, running_sum_quants=0, device_index=0):
-    if isinstance(module, nn.Linear) and not isinstance(module, F16Linear):
-        module.cuda(device_index)
-        module.compile()
-        quantize_(module, int8_weight_only())
-        running_sum_quants += 1
-    elif isinstance(module, F16Linear):
-        module.cuda(device_index)
-    elif isinstance(module, nn.Conv2d):
-        module.cuda(device_index)
-    elif isinstance(module, nn.Embedding):
-        module.cuda(device_index)
-    elif isinstance(module, nn.ConvTranspose2d):
-        module.cuda(device_index)
-    elif isinstance(module, nn.Conv1d):
-        module.cuda(device_index)
-    elif isinstance(module, nn.Conv3d):
-        module.cuda(device_index)
-    elif isinstance(module, nn.ConvTranspose3d):
-        module.cuda(device_index)
-    elif isinstance(module, nn.RMSNorm):
-        module.cuda(device_index)
-    elif isinstance(module, RMSNorm):
-        module.cuda(device_index)
-    elif isinstance(module, nn.LayerNorm):
-        module.cuda(device_index)
-    return running_sum_quants
-
-
-def full_quant(model, max_quants=24, current_quants=0, device_index=0):
-    for module in model.modules():
-        if current_quants < max_quants:
-            current_quants = quant_module(
-                module, current_quants, device_index=device_index
-            )
-    return current_quants
-
-
-@torch.inference_mode()
-def load_pipeline_from_config_path(path: str) -> Model:
-    config = load_config_from_path(path)
-    return load_pipeline_from_config(config)
-
-
-@torch.inference_mode()
-def load_pipeline_from_config(config: ModelSpec) -> Model:
-    models = load_models_from_config(config)
-    config = models.config
-    num_quanted = 0
-    max_quanted = config.num_to_quant
-    flux_device = into_device(config.flux_device)
-    ae_device = into_device(config.ae_device)
-    clip_device = into_device(config.text_enc_device)
-    t5_device = into_device(config.text_enc_device)
-    flux_dtype = into_dtype(config.flow_dtype)
-    device_index = flux_device.index or 0
-    flow_model = models.flow.requires_grad_(False).eval().type(flux_dtype)
-    for block in flow_model.single_blocks:
-        block.cuda(flux_device)
-        if num_quanted < max_quanted:
-            num_quanted = quant_module(
-                block.linear1, num_quanted, device_index=device_index
-            )
-
-    for block in flow_model.double_blocks:
-        block.cuda(flux_device)
-        if num_quanted < max_quanted:
-            num_quanted = full_quant(
-                block, max_quanted, num_quanted, device_index=device_index
-            )
-
-    to_gpu_extras = [
-        "vector_in",
-        "img_in",
-        "txt_in",
-        "time_in",
-        "guidance_in",
-        "final_layer",
-        "pe_embedder",
-    ]
-    for extra in to_gpu_extras:
-        getattr(flow_model, extra).cuda(flux_device).type(flux_dtype)
-    return Model(
-        name=config.version,
-        clip=models.clip,
-        t5=models.t5,
-        model=flow_model,
-        ae=models.ae,
-        dtype=flux_dtype,
-        verbose=False,
-        flux_device=flux_device,
-        ae_device=ae_device,
-        clip_device=clip_device,
-        t5_device=t5_device,
-    )
-
-
-if __name__ == "__main__":
-    pipe = load_pipeline_from_config_path("config-dev.json")
-    o = pipe.generate(
-        prompt="a beautiful asian woman in traditional clothing with golden hairpin and blue eyes, wearing a red kimono with dragon patterns",
-        height=1024,
-        width=1024,
-        seed=13456,
-        num_steps=24,
-        guidance=3.0,
-    )
-    open("out.jpg", "wb").write(o.read())
-    o = pipe.generate(
-        prompt="a beautiful asian woman in traditional clothing with golden hairpin and blue eyes, wearing a red kimono with dragon patterns",
-        height=1024,
-        width=1024,
-        seed=7,
-        num_steps=24,
-        guidance=3.0,
-    )
-    open("out2.jpg", "wb").write(o.read())

flux_pipeline.py

@@ -0,0 +1,461 @@
+import base64
+import io
+import math
+from typing import TYPE_CHECKING, Callable, List
+from PIL import Image
+from einops import rearrange, repeat
+import numpy as np
+
+import torch
+
+from flux_emphasis import get_weighted_text_embeddings_flux
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.benchmark_limit = 20
+torch.set_float32_matmul_precision("high")
+from torch._dynamo import config
+from torch._inductor import config as ind_config
+from pybase64 import standard_b64decode
+
+config.cache_size_limit = 10000000000
+ind_config.force_fuse_int_mm_with_mul = True
+
+from loguru import logger
+from turbojpeg_imgs import TurboImage
+from torchvision.transforms import functional as TF
+from tqdm import tqdm
+from util import (
+    ModelSpec,
+    into_device,
+    into_dtype,
+    load_config_from_path,
+    load_models_from_config,
+)
+
+
+if TYPE_CHECKING:
+    from modules.conditioner import HFEmbedder
+    from modules.flux_model import Flux
+    from modules.autoencoder import AutoEncoder
+
+
+class FluxPipeline:
+    def __init__(
+        self,
+        name: str,
+        offload: bool = False,
+        clip: "HFEmbedder" = None,
+        t5: "HFEmbedder" = None,
+        model: "Flux" = None,
+        ae: "AutoEncoder" = None,
+        dtype: torch.dtype = torch.bfloat16,
+        verbose: bool = False,
+        flux_device: torch.device | str = "cuda:0",
+        ae_device: torch.device | str = "cuda:1",
+        clip_device: torch.device | str = "cuda:1",
+        t5_device: torch.device | str = "cuda:1",
+        config: ModelSpec = None,
+    ):
+
+        self.name = name
+        self.device_flux = (
+            flux_device
+            if isinstance(flux_device, torch.device)
+            else torch.device(flux_device)
+        )
+        self.device_ae = (
+            ae_device
+            if isinstance(ae_device, torch.device)
+            else torch.device(ae_device)
+        )
+        self.device_clip = (
+            clip_device
+            if isinstance(clip_device, torch.device)
+            else torch.device(clip_device)
+        )
+        self.device_t5 = (
+            t5_device
+            if isinstance(t5_device, torch.device)
+            else torch.device(t5_device)
+        )
+        self.dtype = dtype
+        self.offload = offload
+        self.clip: "HFEmbedder" = clip
+        self.t5: "HFEmbedder" = t5
+        self.model: "Flux" = model
+        self.ae: "AutoEncoder" = ae
+        self.rng = torch.Generator(device="cpu")
+        self.turbojpeg = TurboImage()
+        self.verbose = verbose
+        self.ae_dtype = torch.bfloat16
+        self.config = config
+
+    @torch.inference_mode()
+    def prepare(
+        self,
+        img: torch.Tensor,
+        prompt: str | list[str],
+        target_device: torch.device = torch.device("cuda:0"),
+        target_dtype: torch.dtype = torch.float16,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        bs, c, h, w = img.shape
+        if bs == 1 and not isinstance(prompt, str):
+            bs = len(prompt)
+        img = img.unfold(2, 2, 2).unfold(3, 2, 2).permute(0, 2, 3, 1, 4, 5)
+        img = img.reshape(img.shape[0], -1, img.shape[3] * img.shape[4] * img.shape[5])
+        assert img.shape == (
+            bs,
+            (h // 2) * (w // 2),
+            c * 2 * 2,
+        ), f"{img.shape} != {(bs, (h//2)*(w//2), c*2*2)}"
+        if img.shape[0] == 1 and bs > 1:
+            img = img[None].repeat_interleave(bs, dim=0)
+
+        img_ids = torch.zeros(
+            h // 2, w // 2, 3, device=target_device, dtype=target_dtype
+        )
+        img_ids[..., 1] = (
+            img_ids[..., 1]
+            + torch.arange(h // 2, device=target_device, dtype=target_dtype)[:, None]
+        )
+        img_ids[..., 2] = (
+            img_ids[..., 2]
+            + torch.arange(w // 2, device=target_device, dtype=target_dtype)[None, :]
+        )
+
+        img_ids = img_ids[None].repeat(bs, 1, 1, 1).flatten(1, 2)
+        vec, txt, txt_ids = get_weighted_text_embeddings_flux(
+            self,
+            prompt,
+            num_images_per_prompt=bs,
+            device=self.device_clip,
+            target_device=target_device,
+            target_dtype=target_dtype,
+        )
+        return img, img_ids, vec, txt, txt_ids
+
+    @torch.inference_mode()
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def get_lin_function(
+        self, x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
+    ) -> Callable[[float], float]:
+        m = (y2 - y1) / (x2 - x1)
+        b = y1 - m * x1
+        return lambda x: m * x + b
+
+    @torch.inference_mode()
+    def get_schedule(
+        self,
+        num_steps: int,
+        image_seq_len: int,
+        base_shift: float = 0.5,
+        max_shift: float = 1.15,
+        shift: bool = True,
+    ) -> list[float]:
+        # extra step for zero
+        timesteps = torch.linspace(1, 0, num_steps + 1)
+
+        # shifting the schedule to favor high timesteps for higher signal images
+        if shift:
+            # eastimate mu based on linear estimation between two points
+            mu = self.get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+            timesteps = self.time_shift(mu, 1.0, timesteps)
+
+        return timesteps.tolist()
+
+    @torch.inference_mode()
+    def get_noise(
+        self,
+        num_samples: int,
+        height: int,
+        width: int,
+        generator: torch.Generator,
+        dtype=None,
+        device=None,
+    ):
+        if device is None:
+            device = self.device_flux
+        if dtype is None:
+            dtype = self.dtype
+        return torch.randn(
+            num_samples,
+            16,
+            # allow for packing
+            2 * math.ceil(height / 16),
+            2 * math.ceil(width / 16),
+            device=device,
+            dtype=dtype,
+            generator=generator,
+            requires_grad=False,
+        )
+
+    @torch.inference_mode()
+    def into_bytes(self, x: torch.Tensor) -> io.BytesIO:
+        # bring into PIL format and save
+        x = x.clamp(-1, 1)
+        num_images = x.shape[0]
+        images: List[torch.Tensor] = []
+        for i in range(num_images):
+            x = x[i].permute(1, 2, 0).add(1.0).mul(127.5).type(torch.uint8).contiguous()
+            images.append(x)
+        if len(images) == 1:
+            im = images[0]
+        else:
+            im = torch.vstack(images)
+
+        im = self.turbojpeg.encode_torch(im, quality=95)
+        images.clear()
+        return io.BytesIO(im)
+
+    @torch.inference_mode()
+    def vae_decode(self, x: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        x = x.to(self.device_ae)
+        x = self.unpack(x.float(), height, width)
+        with torch.autocast(
+            device_type=self.device_ae.type, dtype=torch.bfloat16, cache_enabled=False
+        ):
+            x = self.ae.decode(x)
+        return x
+
+    def unpack(self, x: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+            h=math.ceil(height / 16),
+            w=math.ceil(width / 16),
+            ph=2,
+            pw=2,
+        )
+
+    @torch.inference_mode()
+    def resize_center_crop(
+        self, img: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        img = TF.resize(img, min(width, height))
+        img = TF.center_crop(img, (height, width))
+        return img
+
+    @torch.inference_mode()
+    def preprocess_latent(
+        self,
+        init_image: torch.Tensor | np.ndarray = None,
+        height: int = 720,
+        width: int = 1024,
+        num_steps: int = 20,
+        strength: float = 1.0,
+        generator: torch.Generator = None,
+        num_images: int = 1,
+    ) -> tuple[torch.Tensor, List[float]]:
+        # prepare input
+
+        if init_image is not None:
+            if isinstance(init_image, np.ndarray):
+                init_image = torch.from_numpy(init_image)
+
+            init_image = (
+                init_image.permute(2, 0, 1)
+                .contiguous()
+                .to(self.device_ae, dtype=self.ae_dtype)
+                .div(127.5)
+                .sub(1)[None, ...]
+            )
+            init_image = self.resize_center_crop(init_image, height, width)
+            with torch.autocast(
+                device_type=self.device_ae.type,
+                dtype=torch.bfloat16,
+                cache_enabled=False,
+            ):
+                init_image = (
+                    self.ae.encode(init_image)
+                    .to(dtype=self.dtype, device=self.device_flux)
+                    .repeat(num_images, 1, 1, 1)
+                )
+
+        x = self.get_noise(
+            num_images,
+            height,
+            width,
+            device=self.device_flux,
+            dtype=self.dtype,
+            generator=generator,
+        )
+        timesteps = self.get_schedule(
+            num_steps=num_steps,
+            image_seq_len=x.shape[-1] * x.shape[-2] // 4,
+            shift=(self.name != "flux-schnell"),
+        )
+        if init_image is not None:
+            t_idx = int((1 - strength) * num_steps)
+            t = timesteps[t_idx]
+            timesteps = timesteps[t_idx:]
+            x = t * x + (1.0 - t) * init_image
+        return x, timesteps
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        prompt: str,
+        width: int = 720,
+        height: int = 1024,
+        num_steps: int = 24,
+        guidance: float = 3.5,
+        seed: int | None = None,
+        init_image: torch.Tensor | str | None = None,
+        strength: float = 1.0,
+        silent: bool = False,
+        num_images: int = 1,
+        return_seed: bool = False,
+    ) -> io.BytesIO:
+        num_steps = 4 if self.name == "flux-schnell" else num_steps
+
+        if isinstance(init_image, str):
+            try:
+                init_image = Image.open(init_image)
+            except Exception as e:
+                init_image = Image.open(io.BytesIO(standard_b64decode(init_image)))
+            init_image = torch.from_numpy(np.array(init_image)).type(torch.uint8)
+
+        # allow for packing and conversion to latent space
+        height = 16 * (height // 16)
+        width = 16 * (width // 16)
+        if isinstance(seed, str):
+            seed = int(seed)
+        if seed is None:
+            seed = self.rng.seed()
+        logger.info(f"Generating with:\nSeed: {seed}\nPrompt: {prompt}")
+
+        generator = torch.Generator(device=self.device_flux).manual_seed(seed)
+        img, timesteps = self.preprocess_latent(
+            init_image=init_image,
+            height=height,
+            width=width,
+            num_steps=num_steps,
+            strength=strength,
+            generator=generator,
+            num_images=num_images,
+        )
+        img, img_ids, vec, txt, txt_ids = self.prepare(
+            img=img,
+            prompt=prompt,
+            target_device=self.device_flux,
+            target_dtype=self.dtype,
+        )
+
+        # this is ignored for schnell
+        guidance_vec = torch.full(
+            (img.shape[0],), guidance, device=self.device_flux, dtype=self.dtype
+        )
+        t_vec = None
+        for t_curr, t_prev in tqdm(
+            zip(timesteps[:-1], timesteps[1:]), total=len(timesteps) - 1, disable=silent
+        ):
+            if t_vec is None:
+                t_vec = torch.full(
+                    (img.shape[0],),
+                    t_curr,
+                    dtype=self.dtype,
+                    device=self.device_flux,
+                )
+            else:
+                t_vec = t_vec.reshape((img.shape[0],)).fill_(t_curr)
+            pred = self.model.forward(
+                img=img,
+                img_ids=img_ids,
+                txt=txt,
+                txt_ids=txt_ids,
+                y=vec,
+                timesteps=t_vec,
+                guidance=guidance_vec,
+            )
+
+            img = img + (t_prev - t_curr) * pred
+
+        torch.cuda.empty_cache()
+
+        # decode latents to pixel space
+        img = self.vae_decode(img, height, width)
+
+        if return_seed:
+            return self.into_bytes(img), seed
+        return self.into_bytes(img)
+
+    @classmethod
+    def load_pipeline_from_config_path(cls, path: str) -> "FluxPipeline":
+        with torch.inference_mode():
+            config = load_config_from_path(path)
+            return cls.load_pipeline_from_config(config)
+
+    @classmethod
+    def load_pipeline_from_config(cls, config: ModelSpec) -> "FluxPipeline":
+        from quantize_swap_and_dispatch import quantize_and_dispatch_to_device
+
+        with torch.inference_mode():
+
+            models = load_models_from_config(config)
+            config = models.config
+            num_layers_to_quantize = config.num_to_quant
+            flux_device = into_device(config.flux_device)
+            ae_device = into_device(config.ae_device)
+            clip_device = into_device(config.text_enc_device)
+            t5_device = into_device(config.text_enc_device)
+            flux_dtype = into_dtype(config.flow_dtype)
+            flow_model = models.flow
+
+            flow_model = quantize_and_dispatch_to_device(
+                flow_model=flow_model,
+                flux_device=flux_device,
+                flux_dtype=flux_dtype,
+                num_layers_to_quantize=num_layers_to_quantize,
+                compile_extras=config.compile_extras,
+                compile_blocks=config.compile_blocks,
+                quantize_extras=config.quantize_extras,
+            )
+
+        return cls(
+            name=config.version,
+            clip=models.clip,
+            t5=models.t5,
+            model=flow_model,
+            ae=models.ae,
+            dtype=flux_dtype,
+            verbose=False,
+            flux_device=flux_device,
+            ae_device=ae_device,
+            clip_device=clip_device,
+            t5_device=t5_device,
+            config=config,
+        )
+
+
+if __name__ == "__main__":
+    pipe = FluxPipeline.load_pipeline_from_config_path(
+        "configs/config-dev-gigaquant.json"
+    )
+    o = pipe.generate(
+        prompt="Street photography portrait of a beautiful asian woman in traditional clothing with golden hairpin and blue eyes, wearing a red kimono with dragon patterns",
+        height=1024,
+        width=1024,
+        num_steps=24,
+        guidance=3.0,
+    )
+    open("out.jpg", "wb").write(o.read())
+    o = pipe.generate(
+        prompt="Street photography portrait of a beautiful asian woman in traditional clothing with golden hairpin and blue eyes, wearing a red kimono with dragon patterns",
+        height=1024,
+        width=1024,
+        num_steps=24,
+        guidance=3.0,
+    )
+    open("out2.jpg", "wb").write(o.read())
+    o = pipe.generate(
+        prompt="Street photography portrait of a beautiful asian woman in traditional clothing with golden hairpin and blue eyes, wearing a red kimono with dragon patterns",
+        height=1024,
+        width=1024,
+        num_steps=24,
+        guidance=3.0,
+    )
+    open("out3.jpg", "wb").write(o.read())

main.py

@@ -1,27 +1,43 @@
 import argparse
 import uvicorn
 from api import app
-from flux_impl import load_pipeline_from_config, load_pipeline_from_config_path
+from flux_pipeline import FluxPipeline
 from util import load_config, ModelVersion
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Launch Flux API server")
     parser.add_argument(
+        "-c",
         "--config-path",
         type=str,
         help="Path to the configuration file, if not provided, the model will be loaded from the command line arguments",
     )
     parser.add_argument(
-        "--port", type=int, default=8088, help="Port to run the server on"
+        "-p",
+        "--port",
+        type=int,
+        default=8088,
+        help="Port to run the server on",
+    )
+    parser.add_argument(
+        "-H",
+        "--host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to run the server on",
+    )
+    parser.add_argument(
+        "-f", "--flow-model-path", type=str, help="Path to the flow model"
+    )
+    parser.add_argument(
+        "-t", "--text-enc-path", type=str, help="Path to the text encoder"
     )
     parser.add_argument(
-        "--host", type=str, default="0.0.0.0", help="Host to run the server on"
+        "-a", "--autoencoder-path", type=str, help="Path to the autoencoder"
     )
-    parser.add_argument("--flow-model-path", type=str, help="Path to the flow model")
-    parser.add_argument("--text-enc-path", type=str, help="Path to the text encoder")
-    parser.add_argument("--autoencoder-path", type=str, help="Path to the autoencoder")
     parser.add_argument(
+        "-m",
         "--model-version",
         type=str,
         choices=["flux-dev", "flux-schnell"],
@@ -29,29 +45,40 @@ def parse_args():
         help="Choose model version",
     )
     parser.add_argument(
+        "-F",
         "--flux-device",
         type=str,
         default="cuda:0",
         help="Device to run the flow model on",
     )
     parser.add_argument(
+        "-T",
         "--text-enc-device",
         type=str,
         default="cuda:0",
         help="Device to run the text encoder on",
     )
     parser.add_argument(
+        "-A",
         "--autoencoder-device",
         type=str,
         default="cuda:0",
         help="Device to run the autoencoder on",
     )
     parser.add_argument(
+        "-q",
         "--num-to-quant",
         type=int,
         default=20,
         help="Number of linear layers in flow transformer (the 'unet') to quantize",
     )
+    parser.add_argument(
+        "-C",
+        "--compile",
+        action="store_true",
+        default=False,
+        help="Compile the flow model with extra optimizations",
+    )
 
     return parser.parse_args()
 
@@ -60,7 +87,7 @@ def main():
     args = parse_args()
 
     if args.config_path:
-        app.state.model = load_pipeline_from_config_path(args.config_path)
+        app.state.model = FluxPipeline.load_pipeline_from_config_path(args.config_path)
     else:
         model_version = (
             ModelVersion.flux_dev
@@ -79,8 +106,10 @@ def main():
             text_enc_dtype="bfloat16",
             ae_dtype="bfloat16",
             num_to_quant=args.num_to_quant,
+            compile_extras=args.compile,
+            compile_blocks=args.compile,
         )
-        app.state.model = load_pipeline_from_config(config)
+        app.state.model = FluxPipeline.load_pipeline_from_config(config)
 
     uvicorn.run(app, host=args.host, port=args.port)
 

main_gr.py

@@ -0,0 +1,132 @@
+import torch
+
+from flux_pipeline import FluxPipeline
+import gradio as gr
+from PIL import Image
+
+
+def create_demo(
+    config_path: str,
+):
+    generator = FluxPipeline.load_pipeline_from_config_path(config_path)
+
+    def generate_image(
+        prompt,
+        width,
+        height,
+        num_steps,
+        guidance,
+        seed,
+        init_image,
+        image2image_strength,
+        add_sampling_metadata,
+    ):
+
+        seed = int(seed)
+        if seed == -1:
+            seed = None
+        out = generator.generate(
+            prompt,
+            width,
+            height,
+            num_steps=num_steps,
+            guidance=guidance,
+            seed=seed,
+            init_image=init_image,
+            strength=image2image_strength,
+            silent=False,
+            num_images=1,
+            return_seed=True,
+        )
+        image_bytes = out[0]
+        return Image.open(image_bytes), str(out[1]), None
+
+    is_schnell = generator.config.version == "flux-schnell"
+
+    with gr.Blocks() as demo:
+        gr.Markdown(f"# Flux Image Generation Demo - Model: {generator.config.version}")
+
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    value='a photo of a forest with mist swirling around the tree trunks. The word "FLUX" is painted over it in big, red brush strokes with visible texture',
+                )
+                do_img2img = gr.Checkbox(
+                    label="Image to Image", value=False, interactive=not is_schnell
+                )
+                init_image = gr.Image(label="Input Image", visible=False)
+                image2image_strength = gr.Slider(
+                    0.0, 1.0, 0.8, step=0.1, label="Noising strength", visible=False
+                )
+
+                with gr.Accordion("Advanced Options", open=False):
+                    width = gr.Slider(128, 8192, 1152, step=16, label="Width")
+                    height = gr.Slider(128, 8192, 640, step=16, label="Height")
+                    num_steps = gr.Slider(
+                        1, 50, 4 if is_schnell else 20, step=1, label="Number of steps"
+                    )
+                    guidance = gr.Slider(
+                        1.0,
+                        10.0,
+                        3.5,
+                        step=0.1,
+                        label="Guidance",
+                        interactive=not is_schnell,
+                    )
+                    seed = gr.Textbox(-1, label="Seed (-1 for random)")
+                    add_sampling_metadata = gr.Checkbox(
+                        label="Add sampling parameters to metadata?", value=True
+                    )
+
+                generate_btn = gr.Button("Generate")
+
+            with gr.Column(min_width="960px"):
+                output_image = gr.Image(label="Generated Image")
+                seed_output = gr.Number(label="Used Seed")
+                warning_text = gr.Textbox(label="Warning", visible=False)
+                # download_btn = gr.File(label="Download full-resolution")
+
+        def update_img2img(do_img2img):
+            return {
+                init_image: gr.update(visible=do_img2img),
+                image2image_strength: gr.update(visible=do_img2img),
+            }
+
+        do_img2img.change(
+            update_img2img, do_img2img, [init_image, image2image_strength]
+        )
+
+        generate_btn.click(
+            fn=generate_image,
+            inputs=[
+                prompt,
+                width,
+                height,
+                num_steps,
+                guidance,
+                seed,
+                init_image,
+                image2image_strength,
+                add_sampling_metadata,
+            ],
+            outputs=[output_image, seed_output, warning_text],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Flux")
+    parser.add_argument(
+        "--config", type=str, default="configs/config-dev.json", help="Config file path"
+    )
+    parser.add_argument(
+        "--share", action="store_true", help="Create a public link to your demo"
+    )
+    args = parser.parse_args()
+
+    demo = create_demo(args.config)
+    demo.launch(share=args.share)

modules/conditioner.py

@@ -2,7 +2,7 @@ from torch import Tensor, nn
 import torch
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 
-from transformers.utils.quantization_config import BitsAndBytesConfig
+from transformers.utils.quantization_config import BitsAndBytesConfig, QuantoConfig
 
 
 class HFEmbedder(nn.Module):
@@ -30,8 +30,8 @@ class HFEmbedder(nn.Module):
                 version,
                 **hf_kwargs,
                 device_map={"": device},
-                quantization_config=BitsAndBytesConfig(
-                    load_in_4bit=True,
+                quantization_config=QuantoConfig(
+                    weights="float8",
                 ),
             )
 

modules/flux_model.py

@@ -1,23 +1,25 @@
+from collections import namedtuple
+import os
 import torch
 
+DISABLE_COMPILE = os.getenv("DISABLE_COMPILE", "0") == "1"
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cudnn.benchmark = True
 torch.backends.cudnn.benchmark_limit = 20
 torch.set_float32_matmul_precision("high")
 import math
-from dataclasses import dataclass
 
-from cublas_linear import CublasLinear as F16Linear
-from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 from torch._dynamo import config
 from torch._inductor import config as ind_config
-from xformers.ops import memory_efficient_attention
+from xformers.ops import memory_efficient_attention_forward
 from pydantic import BaseModel
+from torch.nn import functional as F
 
 config.cache_size_limit = 10000000000
-ind_config.force_fuse_int_mm_with_mul = True
+ind_config.compile_threads = os.cpu_count()
+ind_config.shape_padding = True
 
 
 class FluxParams(BaseModel):
@@ -35,17 +37,16 @@ class FluxParams(BaseModel):
     guidance_embed: bool
 
 
-@torch.compile(mode="reduce-overhead")
+# attention is always same shape each time it's called per H*W, so compile with fullgraph
+@torch.compile(mode="reduce-overhead", fullgraph=True, disable=DISABLE_COMPILE)
 def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
     q, k = apply_rope(q, k, pe)
-    x = memory_efficient_attention(
-        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
-    )
+    x = F.scaled_dot_product_attention(q, k, v).transpose(1, 2)
     x = x.reshape(*x.shape[:-2], -1)
     return x
 
 
-@torch.compile(mode="reduce-overhead")
+@torch.compile(mode="reduce-overhead", disable=DISABLE_COMPILE)
 def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
     scale = torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device) / dim
     omega = 1.0 / (theta**scale)
@@ -119,30 +120,21 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
 class MLPEmbedder(nn.Module):
     def __init__(self, in_dim: int, hidden_dim: int):
         super().__init__()
-        self.in_layer = F16Linear(in_dim, hidden_dim, bias=True)
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
         self.silu = nn.SiLU()
-        self.out_layer = F16Linear(hidden_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
 
     def forward(self, x: Tensor) -> Tensor:
         return self.out_layer(self.silu(self.in_layer(x)))
 
 
-@torch.compile(mode="reduce-overhead", dynamic=True)
-def calculation(
-    x,
-):
-    rrms = torch.rsqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + 1e-6)
-    x = x * rrms
-    return x
-
-
 class RMSNorm(torch.nn.Module):
     def __init__(self, dim: int):
         super().__init__()
         self.scale = nn.Parameter(torch.ones(dim))
 
     def forward(self, x: Tensor):
-        return calculation(x) * self.scale
+        return F.rms_norm(x, self.scale.shape, self.scale, eps=1e-6)
 
 
 class QKNorm(torch.nn.Module):
@@ -163,25 +155,28 @@ class SelfAttention(nn.Module):
         self.num_heads = num_heads
         head_dim = dim // num_heads
 
-        self.qkv = F16Linear(dim, dim * 3, bias=qkv_bias)
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.norm = QKNorm(head_dim)
-        self.proj = F16Linear(dim, dim)
-        self.rearrange = Rearrange("B L (K H D) -> K B H L D", K=3, H=num_heads)
+        self.proj = nn.Linear(dim, dim)
+        self.K = 3
+        self.H = self.num_heads
+        self.KH = self.K * self.H
+
+    def rearrange_for_norm(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        B, L, D = x.shape
+        q, k, v = x.reshape(B, L, self.K, self.H, D // self.KH).permute(2, 0, 3, 1, 4)
+        return q, k, v
 
     def forward(self, x: Tensor, pe: Tensor) -> Tensor:
         qkv = self.qkv(x)
-        q, k, v = self.rearrange(qkv)
+        q, k, v = self.rearrange_for_norm(qkv)
         q, k = self.norm(q, k, v)
         x = attention(q, k, v, pe=pe)
         x = self.proj(x)
         return x
 
 
-@dataclass
-class ModulationOut:
-    shift: Tensor
-    scale: Tensor
-    gate: Tensor
+ModulationOut = namedtuple("ModulationOut", ["shift", "scale", "gate"])
 
 
 class Modulation(nn.Module):
@@ -225,9 +220,9 @@ class DoubleStreamBlock(nn.Module):
 
         self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
         self.img_mlp = nn.Sequential(
-            F16Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
             nn.GELU(approximate="tanh"),
-            F16Linear(mlp_hidden_dim, hidden_size, bias=True),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
         )
 
         self.txt_mod = Modulation(hidden_size, double=True)
@@ -238,13 +233,18 @@ class DoubleStreamBlock(nn.Module):
 
         self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
         self.txt_mlp = nn.Sequential(
-            (F16Linear(hidden_size, mlp_hidden_dim, bias=True)),
+            (nn.Linear(hidden_size, mlp_hidden_dim, bias=True)),
             nn.GELU(approximate="tanh"),
-            (F16Linear(mlp_hidden_dim, hidden_size, bias=True)),
-        )
-        self.rearrange_for_norm = Rearrange(
-            "B L (K H D) -> K B H L D", K=3, H=num_heads
+            (nn.Linear(mlp_hidden_dim, hidden_size, bias=True)),
         )
+        self.K = 3
+        self.H = self.num_heads
+        self.KH = self.K * self.H
+
+    def rearrange_for_norm(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        B, L, D = x.shape
+        q, k, v = x.reshape(B, L, self.K, self.H, D // self.KH).permute(2, 0, 3, 1, 4)
+        return q, k, v
 
     def forward(
         self,
@@ -316,7 +316,7 @@ class SingleStreamBlock(nn.Module):
         # qkv and mlp_in
         self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
         # proj and mlp_out
-        self.linear2 = F16Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
 
         self.norm = QKNorm(head_dim)
 
@@ -325,9 +325,10 @@ class SingleStreamBlock(nn.Module):
 
         self.mlp_act = nn.GELU(approximate="tanh")
         self.modulation = Modulation(hidden_size, double=False)
-        self.rearrange_for_norm = Rearrange(
-            "B L (K H D) -> K B H L D", K=3, H=num_heads
-        )
+
+        self.K = 3
+        self.H = self.num_heads
+        self.KH = self.K * self.H
 
     def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
         mod = self.modulation(vec)[0]
@@ -338,7 +339,8 @@ class SingleStreamBlock(nn.Module):
             [3 * self.hidden_size, self.mlp_hidden_dim],
             dim=-1,
         )
-        q, k, v = self.rearrange_for_norm(qkv)
+        B, L, D = qkv.shape
+        q, k, v = qkv.reshape(B, L, self.K, self.H, D // self.KH).permute(2, 0, 3, 1, 4)
         q, k = self.norm(q, k, v)
         attn = attention(q, k, v, pe=pe)
         output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)).clamp(
@@ -394,7 +396,7 @@ class Flux(nn.Module):
             axes_dim=params.axes_dim,
             dtype=self.dtype,
         )
-        self.img_in = F16Linear(self.in_channels, self.hidden_size, bias=True)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
         self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
         self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
         self.guidance_in = (
@@ -402,7 +404,7 @@ class Flux(nn.Module):
             if params.guidance_embed
             else nn.Identity()
         )
-        self.txt_in = F16Linear(params.context_in_dim, self.hidden_size)
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
 
         self.double_blocks = nn.ModuleList(
             [
@@ -464,10 +466,13 @@ class Flux(nn.Module):
         ids = torch.cat((txt_ids, img_ids), dim=1)
         pe = self.pe_embedder(ids)
 
-        for i, block in enumerate(self.double_blocks):
+        # double stream blocks
+        for block in self.double_blocks:
             img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
 
         img = torch.cat((txt, img), 1)
+
+        # single stream blocks
         for block in self.single_blocks:
             img = block(img, vec=vec, pe=pe)
 
@@ -476,17 +481,14 @@ class Flux(nn.Module):
         return img
 
     @classmethod
-    def from_safetensors(
-        self,
-        model_path: str,
-        model_params: FluxParams,
-        dtype: torch.dtype = torch.bfloat16,
-        device: torch.device = torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        ),
-    ):
+    def from_pretrained(cls, path: str, dtype: torch.dtype = torch.bfloat16) -> "Flux":
+        from util import load_config_from_path
+        from safetensors.torch import load_file
+
+        config = load_config_from_path(path)
+        with torch.device("meta"):
+            klass = cls(params=config.params, dtype=dtype).type(dtype)
 
-        model = Flux(params=model_params, dtype=dtype)
-        model.load_state_dict(model_path.state_dict())
-        model.to(device)
-        return model
+        ckpt = load_file(config.ckpt_path, device="cpu")
+        klass.load_state_dict(ckpt, assign=True)
+        return klass.to("cpu")

quantize_swap_and_dispatch.py

@@ -0,0 +1,241 @@
+from fnmatch import fnmatch
+from typing import List, Optional, Union
+
+import torch
+from click import secho
+from cublas_ops import CublasLinear
+
+from quanto.nn import QModuleMixin, quantize_module, QLinear, QConv2d, QLayerNorm
+from quanto.tensor import Optimizer, qtype, qfloat8
+from torch import nn
+
+
+def _set_module_by_name(parent_module, name, child_module):
+    module_names = name.split(".")
+    if len(module_names) == 1:
+        setattr(parent_module, name, child_module)
+    else:
+        parent_module_name = name[: name.rindex(".")]
+        parent_module = parent_module.get_submodule(parent_module_name)
+        setattr(parent_module, module_names[-1], child_module)
+
+
+def _quantize_submodule(
+    model: torch.nn.Module,
+    name: str,
+    module: torch.nn.Module,
+    weights: Optional[Union[str, qtype]] = None,
+    activations: Optional[Union[str, qtype]] = None,
+    optimizer: Optional[Optimizer] = None,
+):
+    if isinstance(module, CublasLinear):
+        return 0
+    num_quant = 0
+    qmodule = quantize_module(
+        module, weights=weights, activations=activations, optimizer=optimizer
+    )
+    if qmodule is not None:
+        _set_module_by_name(model, name, qmodule)
+        # num_quant += 1
+        qmodule.name = name
+        for name, param in module.named_parameters():
+            # Save device memory by clearing parameters
+            setattr(module, name, None)
+            del param
+        num_quant += 1
+
+    return num_quant
+
+
+def _quantize(
+    model: torch.nn.Module,
+    weights: Optional[Union[str, qtype]] = None,
+    activations: Optional[Union[str, qtype]] = None,
+    optimizer: Optional[Optimizer] = None,
+    include: Optional[Union[str, List[str]]] = None,
+    exclude: Optional[Union[str, List[str]]] = None,
+):
+    """Quantize the specified model submodules
+
+    Recursively quantize the submodules of the specified parent model.
+
+    Only modules that have quantized counterparts will be quantized.
+
+    If include patterns are specified, the submodule name must match one of them.
+
+    If exclude patterns are specified, the submodule must not match one of them.
+
+    Include or exclude patterns are Unix shell-style wildcards which are NOT regular expressions. See
+    https://docs.python.org/3/library/fnmatch.html for more details.
+
+    Note: quantization happens in-place and modifies the original model and its descendants.
+
+    Args:
+        model (`torch.nn.Module`): the model whose submodules will be quantized.
+        weights (`Optional[Union[str, qtype]]`): the qtype for weights quantization.
+        activations (`Optional[Union[str, qtype]]`): the qtype for activations quantization.
+        include (`Optional[Union[str, List[str]]]`):
+            Patterns constituting the allowlist. If provided, module names must match at
+            least one pattern from the allowlist.
+        exclude (`Optional[Union[str, List[str]]]`):
+            Patterns constituting the denylist. If provided, module names must not match
+            any patterns from the denylist.
+    """
+    num_quant = 0
+    if include is not None:
+        include = [include] if isinstance(include, str) else exclude
+    if exclude is not None:
+        exclude = [exclude] if isinstance(exclude, str) else exclude
+    for name, m in model.named_modules():
+        if include is not None and not any(
+            fnmatch(name, pattern) for pattern in include
+        ):
+            continue
+        if exclude is not None and any(fnmatch(name, pattern) for pattern in exclude):
+            continue
+        num_quant += _quantize_submodule(
+            model,
+            name,
+            m,
+            weights=weights,
+            activations=activations,
+            optimizer=optimizer,
+        )
+    return num_quant
+
+
+def _freeze(model):
+    for name, m in model.named_modules():
+        if isinstance(m, QModuleMixin):
+            m.freeze()
+
+
+def _is_block_compilable(module: nn.Module) -> bool:
+    for module in module.modules():
+        if _is_quantized(module):
+            return False
+    if _is_quantized(module):
+        return False
+    return True
+
+
+def _simple_swap_linears(model: nn.Module, root_name: str = ""):
+    for name, module in model.named_children():
+        if _is_linear(module):
+            weights = module.weight.data
+            bias = None
+            if module.bias is not None:
+                bias = module.bias.data
+            with torch.device(module.weight.device):
+                new_cublas = CublasLinear(
+                    module.in_features,
+                    module.out_features,
+                    bias=bias is not None,
+                    device=module.weight.device,
+                    dtype=module.weight.dtype,
+                )
+            new_cublas.weight.data = weights
+            if bias is not None:
+                new_cublas.bias.data = bias
+            setattr(model, name, new_cublas)
+            if root_name == "":
+                secho(f"Replaced {name} with CublasLinear", fg="green")
+            else:
+                secho(f"Replaced {root_name}.{name} with CublasLinear", fg="green")
+        else:
+            if root_name == "":
+                _simple_swap_linears(module, str(name))
+            else:
+                _simple_swap_linears(module, str(root_name) + "." + str(name))
+
+
+def _full_quant(
+    model, max_quants=24, current_quants=0, quantization_dtype: qtype = qfloat8
+):
+    if current_quants < max_quants:
+        current_quants += _quantize(model, quantization_dtype)
+        _freeze(model)
+        print(f"Quantized {current_quants} modules")
+    return current_quants
+
+
+def _is_linear(module: nn.Module) -> bool:
+    return not isinstance(
+        module, (QLinear, QConv2d, QLayerNorm, CublasLinear)
+    ) and isinstance(module, nn.Linear)
+
+
+def _is_quantized(module: nn.Module) -> bool:
+    return isinstance(module, (QLinear, QConv2d, QLayerNorm))
+
+
+def quantize_and_dispatch_to_device(
+    flow_model: nn.Module,
+    flux_device: torch.device = torch.device("cuda"),
+    flux_dtype: torch.dtype = torch.float16,
+    num_layers_to_quantize: int = 20,
+    quantization_dtype: qtype = qfloat8,
+    compile_blocks: bool = True,
+    compile_extras: bool = True,
+    quantize_extras: bool = False,
+):
+    num_quanted = 0
+    flow_model = flow_model.requires_grad_(False).eval().type(flux_dtype)
+    for block in flow_model.single_blocks:
+        block.cuda(flux_device)
+        if num_quanted < num_layers_to_quantize:
+            num_quanted = _full_quant(
+                block,
+                num_layers_to_quantize,
+                num_quanted,
+                quantization_dtype=quantization_dtype,
+            )
+
+    for block in flow_model.double_blocks:
+        block.cuda(flux_device)
+        if num_quanted < num_layers_to_quantize:
+            num_quanted = _full_quant(
+                block,
+                num_layers_to_quantize,
+                num_quanted,
+                quantization_dtype=quantization_dtype,
+            )
+
+    to_gpu_extras = [
+        "vector_in",
+        "img_in",
+        "txt_in",
+        "time_in",
+        "guidance_in",
+        "final_layer",
+        "pe_embedder",
+    ]
+
+    if compile_blocks:
+        for i, block in enumerate(flow_model.single_blocks):
+            if _is_block_compilable(block):
+                block.compile()
+                secho(f"Compiled block {i}", fg="green")
+        for i, block in enumerate(flow_model.double_blocks):
+            if _is_block_compilable(block):
+                block.compile()
+                secho(f"Compiled block {i}", fg="green")
+
+    _simple_swap_linears(flow_model)
+    for extra in to_gpu_extras:
+        m_extra = getattr(flow_model, extra).cuda(flux_device).type(flux_dtype)
+        if compile_blocks:
+            if extra in ["time_in", "vector_in", "guidance_in", "final_layer"]:
+                m_extra.compile()
+                secho(
+                    f"Compiled extra {extra} -- {m_extra.__class__.__name__}",
+                    fg="green",
+                )
+        elif quantize_extras:
+            _full_quant(
+                m_extra,
+                current_quants=num_quanted,
+                max_quants=num_layers_to_quantize,
+                quantization_dtype=quantization_dtype,
+            )
+    return flow_model

sampling.py

@@ -32,7 +32,12 @@ def get_noise(
 
 @torch.inference_mode()
 def prepare(
-    t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]
+    t5: HFEmbedder,
+    clip: HFEmbedder,
+    img: Tensor,
+    prompt: str | list[str],
+    target_device: torch.device = torch.device("cuda:0"),
+    target_dtype: torch.dtype = torch.float16,
 ) -> dict[str, Tensor]:
     bs, c, h, w = img.shape
     if bs == 1 and not isinstance(prompt, str):
@@ -42,28 +47,34 @@ def prepare(
     if img.shape[0] == 1 and bs > 1:
         img = repeat(img, "1 ... -> bs ...", bs=bs)
 
-    img_ids = torch.zeros(h // 2, w // 2, 3)
-    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
-    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = torch.zeros(h // 2, w // 2, 3, device=target_device, dtype=target_dtype)
+    img_ids[..., 1] = (
+        img_ids[..., 1]
+        + torch.arange(h // 2, device=target_device, dtype=target_dtype)[:, None]
+    )
+    img_ids[..., 2] = (
+        img_ids[..., 2]
+        + torch.arange(w // 2, device=target_device, dtype=target_dtype)[None, :]
+    )
     img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
 
     if isinstance(prompt, str):
         prompt = [prompt]
-    txt = t5(prompt)
+    txt = t5(prompt).to(target_device, dtype=target_dtype)
     if txt.shape[0] == 1 and bs > 1:
         txt = repeat(txt, "1 ... -> bs ...", bs=bs)
-    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3, device=target_device, dtype=target_dtype)
 
-    vec = clip(prompt)
+    vec = clip(prompt).to(target_device, dtype=target_dtype)
     if vec.shape[0] == 1 and bs > 1:
         vec = repeat(vec, "1 ... -> bs ...", bs=bs)
 
     return {
         "img": img,
-        "img_ids": img_ids.to(img.device),
-        "txt": txt.to(img.device),
-        "txt_ids": txt_ids.to(img.device),
-        "vec": vec.to(img.device),
+        "img_ids": img_ids,
+        "txt": txt,
+        "txt_ids": txt_ids,
+        "vec": vec,
     }
 
 
@@ -116,11 +127,6 @@ def denoise(
     from tqdm import tqdm
 
     # this is ignored for schnell
-    img = img.to(device=device, dtype=dtype)
-    img_ids = img_ids.to(device=device, dtype=dtype)
-    txt = txt.to(device=device, dtype=dtype)
-    txt_ids = txt_ids.to(device=device, dtype=dtype)
-    vec = vec.to(device=device, dtype=dtype)
     guidance_vec = torch.full((img.shape[0],), guidance, device=device, dtype=dtype)
     for t_curr, t_prev in tqdm(
         zip(timesteps[:-1], timesteps[1:]), total=len(timesteps) - 1

util.py

@@ -36,6 +36,9 @@ class ModelSpec(BaseModel):
     ae_dtype: str = "bfloat16"
     text_enc_dtype: str = "bfloat16"
     num_to_quant: Optional[int] = 20
+    quantize_extras: bool = False
+    compile_extras: bool = False
+    compile_blocks: bool = False
 
     model_config: ConfigDict = {
         "arbitrary_types_allowed": True,
@@ -93,6 +96,8 @@ def load_config(
     ae_dtype: str = "bfloat16",
     text_enc_dtype: str = "bfloat16",
     num_to_quant: Optional[int] = 20,
+    compile_extras: bool = False,
+    compile_blocks: bool = False,
 ):
     text_enc_device = str(parse_device(text_enc_device))
     ae_device = str(parse_device(ae_device))
@@ -144,6 +149,8 @@ def load_config(
         text_enc_dtype=text_enc_dtype,
         text_enc_max_length=512 if name == ModelVersion.flux_dev else 256,
         num_to_quant=num_to_quant,
+        compile_extras=compile_extras,
+        compile_blocks=compile_blocks,
     )