Make Qwen 8B work with TextGenerate node. (Comfy-Org#13160)

Author: comfyanonymous

comfy/ops.py

@@ -928,6 +928,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                 weight = state_dict.pop(weight_key, None)
                 if weight is None:
                     logging.warning(f"Missing weight for layer {layer_name}")
+                    self.weight = None
                     return
 
                 manually_loaded_keys = [weight_key]
@@ -1034,6 +1035,9 @@ def state_dict(self, *args, destination=None, prefix="", **kwargs):
                 if self.bias is not None:
                     sd["{}bias".format(prefix)] = self.bias
 
+                if self.weight is None:
+                    return sd
+
                 if isinstance(self.weight, QuantizedTensor):
                     sd_out = self.weight.state_dict("{}weight".format(prefix))
                     for k in sd_out:

comfy/text_encoders/llama.py

@@ -224,7 +224,7 @@ class Qwen3_8BConfig:
     k_norm = "gemma3"
     rope_scale = None
     final_norm: bool = True
-    lm_head: bool = False
+    lm_head: bool = True
     stop_tokens = [151643, 151645]
 
 @dataclass
@@ -912,6 +912,9 @@ def sample_token(self, logits, temperature, top_k, top_p, min_p, repetition_pena
 class BaseQwen3:
     def logits(self, x):
         input = x[:, -1:]
+        if self.model.config.lm_head:
+            return self.model.lm_head(input)
+
         module = self.model.embed_tokens
 
         offload_stream = None