From 3d673c5c098ecaa6e6027f834659ba8de534ca32 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 26 Apr 2025 21:16:43 +0200
Subject: [PATCH] Wan FLF2V convert logic

Forgot to commit this, also fixes print -> logging
---
 tools/convert.py | 11 +++++++----
 tools/lcpp.patch | 16 +++++++++++-----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/tools/convert.py b/tools/convert.py
index 63cc713..fa12655 100644
--- a/tools/convert.py
+++ b/tools/convert.py
@@ -2,6 +2,7 @@
 import os
 import gguf
 import torch
+import logging
 import argparse
 from tqdm import tqdm
 from safetensors.torch import load_file, save_file
@@ -54,7 +55,9 @@ class ModelHiDream(ModelTemplate):
         )
     ]
     keys_hiprec = [
-        ".ff_i.gate.weight" # nn.parameter, can't load from BF16 ver
+        # nn.parameter, can't load from BF16 ver
+        ".ff_i.gate.weight",
+        "img_emb.emb_pos"
     ]
 
 class ModelHyVid(ModelTemplate):
@@ -275,7 +278,7 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False):
     # load & run model detection logic
     state_dict = load_state_dict(path)
     model_arch = detect_arch(state_dict)
-    print(f"* Architecture detected from input: {model_arch.arch}")
+    logging.info(f"* Architecture detected from input: {model_arch.arch}")
 
     # detect & set dtype for output file
     dtypes = [x.dtype for x in state_dict.values()]
@@ -317,8 +320,8 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False):
 
     fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors"
     if os.path.isfile(fix):
-        print(f"\n### Warning! Fix file found at '{fix}'")
-        print(f" you most likely need to run 'fix_5d_tensors.py' after quantization.")
+        logging.warning(f"\n### Warning! Fix file found at '{fix}'")
+        logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.")
 
     return dst_path, model_arch
 
diff --git a/tools/lcpp.patch b/tools/lcpp.patch
index 17cba8d..cecb849 100644
--- a/tools/lcpp.patch
+++ b/tools/lcpp.patch
@@ -39,7 +39,7 @@ index b16c462f..6d1568f1 100644
      const int idx = gguf_find_tensor(ctx, name);
      if (idx < 0) {
 diff --git a/src/llama.cpp b/src/llama.cpp
-index 24e1f1f0..5e944837 100644
+index 24e1f1f0..39045ca5 100644
 --- a/src/llama.cpp
 +++ b/src/llama.cpp
 @@ -205,6 +205,15 @@ enum llm_arch {
@@ -255,7 +255,7 @@ index 24e1f1f0..5e944837 100644
      }
  
      size_t total_size_org = 0;
-@@ -18547,6 +18716,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18547,6 +18716,51 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              ctx_outs[i_split] = gguf_init_empty();
          }
          gguf_add_tensor(ctx_outs[i_split], tensor);
@@ -297,11 +297,17 @@ index 24e1f1f0..5e944837 100644
 +                gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim);
 +                LLAMA_LOG_INFO("\n%s: Correcting shape for Wan: [key:%s]\n", __func__, tensor->name);
 +            }
++            // FLF2V model only
++            if (name == "img_emb.emb_pos") {
++                const int n_dim = 3;
++                gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim);
++                LLAMA_LOG_INFO("\n%s: Correcting shape for Wan FLF2V: [key:%s]\n", __func__, tensor->name);
++            }
 +        }
      }
  
      // Set split info if needed
-@@ -18647,6 +18855,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18647,6 +18861,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
          // do not quantize relative position bias (T5)
          quantize &= name.find("attn_rel_b.weight") == std::string::npos;
  
@@ -394,7 +400,7 @@ index 24e1f1f0..5e944837 100644
          enum ggml_type new_type;
          void * new_data;
          size_t new_size;
-@@ -18655,6 +18949,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18655,6 +18955,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              new_type = default_type;
  
              // get more optimal quantization type based on the tensor shape, layer, etc.
@@ -404,7 +410,7 @@ index 24e1f1f0..5e944837 100644
              if (!params->pure && ggml_is_quantized(default_type)) {
                  new_type = llama_tensor_get_type(qs, new_type, tensor, ftype);
              }
-@@ -18664,6 +18961,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18664,6 +18967,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) {
                  new_type = params->output_tensor_type;
              }