Update src/pipeline.py

src/pipeline.py

@@ -50,7 +50,7 @@ class W8A16LinearLayer(nn.Module):
 def replace_linear_with_target_and_quantize(module, target_class, module_name_to_exclude):
     # with open("/root/.cache/huggingface/hub/output_layers.txt", "a") as f:
     for name, child in module.named_children():
-        if isinstance(child, nn.Linear) and ( 'to_q' in name or 'to_k' in name or 'to_v' in name ): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or
+        if isinstance(child, nn.Linear) and (  'add_k_proj' in name or 'add_v_proj' in name or 'add_q_proj' in name ): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or
             old_bias = child.bias
             old_weight = child.weight
             new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype)
@@ -98,7 +98,8 @@ def load_pipeline() -> Pipeline:
     pipeline.text_encoder.to(memory_format=torch.channels_last)
     pipeline.transformer.to(memory_format=torch.channels_last)
     replace_linear_with_target_and_quantize(pipeline.transformer, W8A16LinearLayer, [])
-    # exit()    
+    pipeline.transformer.save_pretrained("/root/.cache/huggingface/hub/transformer-flux")
+    exit()    
 
     pipeline.vae.to(memory_format=torch.channels_last)
     pipeline.vae = torch.compile(pipeline.vae)