fix(flux): support .weight variant for attention norm parameter mappings

[Fortyseven]

Summary

Some Flux model checkpoints use *.weight instead of *.scale for attention normalization layer parameters (query_norm and key_norm). The model loader's key mapping functions only checked for .scale, causing these checkpoints to fail loading with unmapped keys.

This change updates all affected mapping functions across multiple Flux model type handlers to also match the .weight variant, mapping it to the same internal parameter name as .scale.

Related Issues / Discussions

https://discord.com/channels/1020123559063990373/1083864753543331981/1504216036168171570

QA Instructions

This was tested with Flux2-Klein-9B-True-V2. (Q4_K specifically, but other sizes exhibit the issue too.)

Installation goes fine, and it's recognized. However, attempting a generation results in:

Server Error

RuntimeError: Error(s) in loading state_dict for Flux2Transformer2DModel: Missing key(s) in state_dict: "transformer_blocks.0.attn.norm_q.weight", "transformer_blocks.0.attn.norm_k.weight",...

After this patch: image generation works.

Note

It's important to note that while I have quite a long history in software development, I am on the very fringes of this particular tech. So I may be overlooking something in this 'fix' that's just outside my understanding. ;)