Update src/pipeline.py

src/pipeline.py

@@ -28,20 +28,19 @@ REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
 Pipeline = None
 
 
-# ---------------- NF4 ----------------
-def functional_linear_4bits(x, weight, bias):
-    out = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
-    out = out.to(x)
-    return out
 
+def quantized_matrix_multiply(x, weight, bias):
+    """Perform matrix multiplication for 4-bit quantized weights."""
+    output = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
+    return output.to(x)
 
-def quant_state_copier(state, device=None):
+def copy_quant_state(state, device=None):
+    """Create a copy of quantization state for a given device."""
     if state is None:
         return None
 
     device = device or state.absmax.device
-
-    state2 = (
+    nested_state = (
         QuantState(
             absmax=state.state2.absmax.to(device),
             shape=state.state2.shape,
@@ -50,8 +49,7 @@ def quant_state_copier(state, device=None):
             quant_type=state.state2.quant_type,
             dtype=state.state2.dtype,
         )
-        if state.nested
-        else None
+        if state.nested else None
     )
 
     return QuantState(
@@ -62,100 +60,45 @@ def quant_state_copier(state, device=None):
         quant_type=state.quant_type,
         dtype=state.dtype,
         offset=state.offset.to(device) if state.nested else None,
-        state2=state2,
+        state2=nested_state,
     )
 
-
-class Forge_Params_4Bit(Params4bit):
+class QuantizedModelParams(Params4bit):
     def to(self, *args, **kwargs):
-        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        device, dtype, non_blocking, _ = 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)
-        else:
-            n = Forge_Params_4Bit(
-                torch.nn.Parameter.to(self, device=device, dtype=dtype, non_blocking=non_blocking),
-                requires_grad=self.requires_grad,
-                quant_state=quant_state_copier(self.quant_state, device),
-                compress_statistics=False,
-                blocksize=64,
-                quant_type=self.quant_type,
-                quant_storage=self.quant_storage,
-                bnb_quantized=self.bnb_quantized,
-                module=self.module
-            )
-            self.module.quant_state = n.quant_state
-            self.data = n.data
-            self.quant_state = n.quant_state
-            return n
-
-
-class Force_Loader_4Bits(torch.nn.Module):
-    def __init__(self, *, device, dtype, quant_type, **kwargs):
+        
+        updated_params = QuantizedModelParams(
+            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),
+            compress_statistics=False,
+            blocksize=64,
+            quant_type=self.quant_type,
+            quant_storage=self.quant_storage,
+            bnb_quantized=self.bnb_quantized,
+            module=self.module
+        )
+        self.module.quant_state = updated_params.quant_state
+        self.data = updated_params.data
+        self.quant_state = updated_params.quant_state
+        return updated_params
+
+class QuantizedLinearLayer(torch.nn.Module):
+    def __init__(self, *args, device=None, dtype=None, **kwargs):
         super().__init__()
         self.dummy = torch.nn.Parameter(torch.empty(1, device=device, dtype=dtype))
         self.weight = None
         self.quant_state = None
         self.bias = None
-        self.quant_type = quant_type
-
-    def _save_to_state_dict(self, destination, prefix, keep_vars):
-        super()._save_to_state_dict(destination, prefix, keep_vars)
-        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.")}
-
-        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 = Forge_Params_4Bit.from_prequantized(
-                data=state_dict[prefix + 'weight'],
-                quantized_stats=quant_state_dict,
-                requires_grad=False,
-                device=torch.device('cuda'),
-                module=self
-            )
-            self.quant_state = self.weight.quant_state
-
-            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:
-                self.weight = Forge_Params_4Bit(
-                    state_dict[prefix + 'weight'].to(self.dummy),
-                    requires_grad=False,
-                    compress_statistics=True,
-                    quant_type=self.quant_type,
-                    quant_storage=torch.uint8,
-                    module=self,
-                )
-                self.quant_state = self.weight.quant_state
-
-            if prefix + 'bias' in state_dict:
-                self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
-
-            del self.dummy
-        else:
-            super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
-
-
-class CustomLinear(Force_Loader_4Bits):
-    def __init__(self, *args, device=None, dtype=None, **kwargs):
-        super().__init__(device=device, dtype=dtype, quant_type='nf4')
+        self.quant_type = 'nf4'
 
     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)
+        return quantized_matrix_multiply(x, self.weight, self.bias)
 
 
 class InitModel:
@@ -170,16 +113,6 @@ class InitModel:
         )
         return text_encoder.to(memory_format=torch.channels_last)
 
-    @staticmethod
-    def load_vae() -> AutoencoderTiny:
-        print("Loading VAE model...")
-        vae = AutoencoderTiny.from_pretrained(
-            "XiangquiAI/FLUX_Vae_Model",
-            revision="103bcc03998f48ef311c100ee119f1b9942132ab",
-            torch_dtype=torch.bfloat16,
-        )
-        return vae
-
     @staticmethod
     def load_transformer(trans_path: str) -> FluxTransformer2DModel:
         print("Loading transformer model...")
@@ -194,8 +127,7 @@ class InitModel:
 
 def load_pipeline() -> Pipeline:
 
-    t5_encoder_2 = InitModel.load_text_encoder()
-    vae = InitModel.load_vae()
+
 
     transformer_path = os.path.join(HF_HUB_CACHE, "models--MyApricity--Flux_Transformer_float8/snapshots/66c5f182385555a00ec90272ab711bb6d3c197db")
     transformer = InitModel.load_transformer(transformer_path)
@@ -203,9 +135,7 @@ def load_pipeline() -> Pipeline:
 
     pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
                         revision=REVISION, 
-                        vae=vae, 
                         transformer=transformer, 
-                        text_encoder_2=t5_encoder_2, 
                         torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
 
@@ -213,19 +143,25 @@ def load_pipeline() -> Pipeline:
         # Enable some options for better vae
         pipeline.enable_vae_slicing()
         pipeline.enable_vae_tiling()
-        torch.nn.LinearLayer = CustomLinear
+        torch.nn.LinearLayer = QuantizedLinearLayer
     except:
         print("Debug here")
+
+    try:
+        pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
+
+    except:
+        print("nothing")
         
 
-    prms = [
-        "melanogen, tiptilt",
+    ps = [
+        "overgross, mandative, inventful, braunite, penneeck",
         "melanogen, endosome, apical, polymyodous, ",
         "buffer, cutie, buttinsky, prototrophic",
         "puzzlehead",
     ]
     
-    for warmprompt in prms:
+    for warmprompt in ps:
         pipeline(prompt=warmprompt, 
                         width=1024,
                         height=1024,