Update app.py

app.py

@@ -66,42 +66,15 @@ ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="va
 # ---------------- NF4 ----------------
 
 def functional_linear_4bits(x, weight, bias):
+    import bitsandbytes as bnb
     out = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
     out = out.to(x)
     return out
 
-def copy_quant_state(state: QuantState, device: torch.device = None) -> QuantState:
-    if state is None:
-        return None
-
-    device = device or state.absmax.device
-
-    state2 = (
-        QuantState(
-            absmax=state.state2.absmax.to(device),
-            shape=state.state2.shape,
-            code=state.state2.code.to(device),
-            blocksize=state.state2.blocksize,
-            quant_type=state.state2.quant_type,
-            dtype=state.state2.dtype,
-        )
-        if state.nested
-        else None
-    )
-
-    return QuantState(
-        absmax=state.absmax.to(device),
-        shape=state.shape,
-        code=state.code.to(device),
-        blocksize=state.blocksize,
-        quant_type=state.quant_type,
-        dtype=state.dtype,
-        offset=state.offset.to(device) if state.nested else None,
-        state2=state2,
-    )
-
 class ForgeParams4bit(Params4bit):
+    """Subclass to force re-quantization to GPU if needed."""
     def to(self, *args, **kwargs):
+        import torch
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
         if device is not None and device.type == "cuda" and not self.bnb_quantized:
             return self._quantize(device)
@@ -109,7 +82,7 @@ class ForgeParams4bit(Params4bit):
             n = ForgeParams4bit(
                 torch.nn.Parameter.to(self, device=device, dtype=dtype, non_blocking=non_blocking),
                 requires_grad=self.requires_grad,
-                quant_state=copy_quant_state(self.quant_state, device),
+                quant_state=self.quant_state,
                 compress_statistics=False,
                 blocksize=64,
                 quant_type=self.quant_type,
@@ -122,10 +95,10 @@ class ForgeParams4bit(Params4bit):
             self.quant_state = n.quant_state
             return n
 
-class ForgeLoader4Bit(torch.nn.Module):
+class ForgeLoader4Bit(nn.Module):
     def __init__(self, *, device, dtype, quant_type, **kwargs):
         super().__init__()
-        self.dummy = torch.nn.Parameter(torch.empty(1, device=device, dtype=dtype))
+        self.dummy = nn.Parameter(torch.empty(1, device=device, dtype=dtype))
         self.weight = None
         self.quant_state = None
         self.bias = None
@@ -133,18 +106,22 @@ class ForgeLoader4Bit(torch.nn.Module):
 
     def _save_to_state_dict(self, destination, prefix, keep_vars):
         super()._save_to_state_dict(destination, prefix, keep_vars)
+        from bitsandbytes.nn.modules import QuantState
         quant_state = getattr(self.weight, "quant_state", None)
         if quant_state is not None:
             for k, v in quant_state.as_dict(packed=True).items():
                 destination[prefix + "weight." + k] = v if keep_vars else v.detach()
         return
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
-        quant_state_keys = {k[len(prefix + "weight."):] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        from bitsandbytes.nn.modules import Params4bit
+        import torch
 
+        quant_state_keys = {k[len(prefix + "weight."):] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
         if any('bitsandbytes' in k for k in quant_state_keys):
             quant_state_dict = {k: state_dict[prefix + "weight." + k] for k in quant_state_keys}
-
             self.weight = ForgeParams4bit.from_prequantized(
                 data=state_dict[prefix + 'weight'],
                 quantized_stats=quant_state_dict,
@@ -156,7 +133,6 @@ class ForgeLoader4Bit(torch.nn.Module):
 
             if prefix + 'bias' in state_dict:
                 self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
-
             del self.dummy
         elif hasattr(self, 'dummy'):
             if prefix + 'weight' in state_dict:
@@ -183,12 +159,11 @@ class Linear(ForgeLoader4Bit):
 
     def forward(self, x):
         self.weight.quant_state = self.quant_state
-
         if self.bias is not None and self.bias.dtype != x.dtype:
             self.bias.data = self.bias.data.to(x.dtype)
-
         return functional_linear_4bits(x, self.weight, self.bias)
 
+import torch.nn as nn
 nn.Linear = Linear
 
 # ---------------- Model ----------------
@@ -200,6 +175,7 @@ def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
     return x
 
 def rope(pos, dim, theta):
+    import torch
     scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
     omega = 1.0 / (theta ** scale)
     out = pos.unsqueeze(-1) * omega.unsqueeze(0)
@@ -225,6 +201,7 @@ class EmbedND(nn.Module):
         self.axes_dim = axes_dim
 
     def forward(self, ids: Tensor) -> Tensor:
+        import torch
         n_axes = ids.shape[-1]
         emb = torch.cat(
             [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
@@ -233,6 +210,7 @@ class EmbedND(nn.Module):
         return emb.unsqueeze(1)
 
 def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    import torch, math
     t = time_factor * t
     half = dim // 2
     freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
@@ -260,6 +238,7 @@ class RMSNorm(torch.nn.Module):
         self.scale = nn.Parameter(torch.ones(dim))
 
     def forward(self, x: Tensor):
+        import torch
         x_dtype = x.dtype
         x = x.float()
         rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
@@ -295,6 +274,8 @@ class SelfAttention(nn.Module):
         x = self.proj(x)
         return x
 
+from dataclasses import dataclass
+
 @dataclass
 class ModulationOut:
     shift: Tensor
@@ -308,12 +289,11 @@ class Modulation(nn.Module):
         self.multiplier = 6 if double else 3
         self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
 
-    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+    def forward(self, vec: Tensor):
         out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
-        return (
-            ModulationOut(*out[:3]),
-            ModulationOut(*out[3:]) if self.is_double else None,
-        )
+        first = ModulationOut(*out[:3])
+        second = ModulationOut(*out[3:]) if self.is_double else None
+        return first, second
 
 class DoubleStreamBlock(nn.Module):
     def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
@@ -424,6 +404,8 @@ class LastLayer(nn.Module):
         x = self.linear(x)
         return x
 
+from dataclasses import dataclass, field
+
 @dataclass
 class FluxParams:
     in_channels: int = 64
@@ -516,6 +498,7 @@ class Flux(nn.Module):
         return img
 
 def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
+    import torch
     bs, c, h, w = img.shape
     if bs == 1 and not isinstance(prompt, str):
         bs = len(prompt)
@@ -544,11 +527,13 @@ def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[st
     }
 
 def time_shift(mu: float, sigma: float, t: Tensor):
+    import math
     return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
 
 def get_lin_function(
     x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
 ) -> Callable[[float], float]:
+    import math
     m = (y2 - y1) / (x2 - x1)
     b = y1 - m * x1
     return lambda x: m * x + b
@@ -560,6 +545,8 @@ def get_schedule(
     max_shift: float = 1.15,
     shift: bool = True,
 ) -> list[float]:
+    import torch
+    import math
     timesteps = torch.linspace(1, 0, num_steps + 1)
     if shift:
         mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
@@ -567,7 +554,7 @@ def get_schedule(
     return timesteps.tolist()
 
 def denoise(
-    model: "Flux",
+    model: Flux,
     img: Tensor,
     img_ids: Tensor,
     txt: Tensor,
@@ -576,6 +563,7 @@ def denoise(
     timesteps: list[float],
     guidance: float = 4.0,
 ):
+    import torch
     guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
     for t_curr, t_prev in tqdm(zip(timesteps[:-1], timesteps[1:]), total=len(timesteps) - 1):
         t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
@@ -620,6 +608,7 @@ def get_image(image) -> torch.Tensor | None:
     img: torch.Tensor = transform(image)
     return img[None, ...]
 
+# Load the NF4 quantized checkpoint
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
@@ -637,16 +626,16 @@ def generate_image(
     progress=gr.Progress(track_tqdm=True),
 ):
     if seed == 0:
-        seed = int(random.random() * 1000000)
-        
+        seed = int(random.random() * 1_000_000)
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch_device = torch.device(device)
-    
+
     global model, model_zero_init
     if not model_zero_init:
         model = model.to(torch_device)
         model_zero_init = True
-    
+
     if do_img2img and init_image is not None:
         init_image = get_image(init_image)
         if resize_img:
@@ -670,11 +659,10 @@ def generate_image(
         generator=generator
     )
 
-    num_steps = inference_steps
-    timesteps = get_schedule(num_steps, (x.shape[-1] * x.shape[-2]) // 4, shift=True)
+    timesteps = get_schedule(inference_steps, (x.shape[-1] * x.shape[-2]) // 4, shift=True)
 
     if do_img2img and init_image is not None:
-        t_idx = int((1 - image2image_strength) * num_steps)
+        t_idx = int((1 - image2image_strength) * inference_steps)
         t = timesteps[t_idx]
         timesteps = timesteps[t_idx:]
         x = t * x + (1.0 - t) * init_image.to(x.dtype)
@@ -682,9 +670,11 @@ def generate_image(
     inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
     x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
     x = unpack(x.float(), height, width)
+
     with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
         x = (x / ae.config.scaling_factor) + ae.config.shift_factor
         x = ae.decode(x).sample
+
     x = x.clamp(-1, 1)
     x = rearrange(x[0], "c h w -> h w c")
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
@@ -717,7 +707,14 @@ def create_demo():
                 seed = gr.Number(label="Seed", precision=-1)
                 do_img2img = gr.Checkbox(label="Image to Image", value=False)
                 init_image = gr.Image(label="Initial Image", visible=False)
-                image2image_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Noising Strength", value=0.8, visible=False)
+                image2image_strength = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    label="Noising Strength",
+                    value=0.8,
+                    visible=False
+                )
                 resize_img = gr.Checkbox(label="Resize Initial Image", value=True, visible=False)
                 generate_button = gr.Button("Generate", variant="primary")
             with gr.Column():
@@ -732,11 +729,20 @@ def create_demo():
 
         generate_button.click(
             fn=generate_image,
-            inputs=[prompt, width, height, guidance, inference_steps, seed, do_img2img, init_image, image2image_strength, resize_img],
+            inputs=[
+                prompt, width, height, guidance,
+                inference_steps, seed, do_img2img,
+                init_image, image2image_strength, resize_img
+            ],
             outputs=[output_image, output_seed]
         )
     return demo
 
 if __name__ == "__main__":
+    # Create the demo
     demo = create_demo()
-    demo.launch()
+    # Enable the queue to handle concurrency
+    demo.queue()
+    # Launch with show_api=False and share=True to avoid the "bool is not iterable" error
+    # and the "ValueError: When localhost is not accessible..." error.
+    demo.launch(show_api=False, share=True, server_name="0.0.0.0")