From 0fd03c75521c7f6e947a9f04e01dacea61b5a8a2 Mon Sep 17 00:00:00 2001 From: elitexp Date: Wed, 4 Sep 2024 21:52:51 +0545 Subject: [PATCH] @Xuzzo - Weight Adjustment as Original Transformer Co-authored-by: Fabio Peruzzo <39661245+Xuzzo@users.noreply.github.com> --- src/flux_1/weights/weight_handler.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 863559e..0369558 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -105,6 +105,16 @@ class WeightHandler: weights = [WeightHandler._reshape_weights(k, v) for k, v in weights] weights = WeightHandler._flatten(weights) unflatten = tree_unflatten(weights) + for block in unflatten["transformer"]["transformer_blocks"]: + block["ff"] = { + "linear1": block["ff"]["net"][0]["proj"], + "linear2": block["ff"]["net"][2] + } + if block.get("ff_context") is not None: + block["ff_context"] = { + "linear1": block["ff_context"]["net"][0]["proj"], + "linear2": block["ff_context"]["net"][2] + } return unflatten, quantization_level