diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py
index f9b2862a..3a5387a8 100644
--- a/zeromodels/__init__.py
+++ b/zeromodels/__init__.py
@@ -8,7 +8,7 @@
 )
 from zeromodels.version import version
 
-__version__ = "1.3.4"
+__version__ = "1.3.5"
 
 __all__ = [
     "auto",
diff --git a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py
index eeafd2c2..d7ee8d5b 100644
--- a/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py
+++ b/zeromodels/models/qwen_image/convert_qwen_image_diffusers_to_keras.py
@@ -117,7 +117,6 @@ def load_json(filename, subfolder=None):
     )
 
 
-
 def transfer_qwen_image(
     repo, token=None, dtype="float16", build_sample_size=16, config=None
 ):
@@ -140,9 +139,7 @@ def transfer_qwen_image(
         model = QwenImageModel(**flat)
 
     vae_mapping = {
-        k: v
-        for k, v in WEIGHT_NAME_MAPPING.items()
-        if k not in ("/gamma", "gamma")
+        k: v for k, v in WEIGHT_NAME_MAPPING.items() if k not in ("/gamma", "gamma")
     }
     for step, (component, subfolder, mapping, index_name, filename) in enumerate(
         (
@@ -176,25 +173,29 @@ def transfer_qwen_image(
             }
 
             class _State(dict):
+                def __init__(self, weight_map, shard_paths):
+                    self.weight_map = weight_map
+                    self.shard_paths = shard_paths
+
                 def __contains__(self, key):
-                    return key in weight_map
+                    return key in self.weight_map
 
                 def __getitem__(self, key):
                     with safe_open(
-                        shard_paths[weight_map[key]], framework="np"
+                        self.shard_paths[self.weight_map[key]], framework="np"
                     ) as shard:
                         return shard.get_tensor(key)
 
                 def keys(self):
-                    return weight_map.keys()
+                    return self.weight_map.keys()
 
                 def __iter__(self):
-                    return iter(weight_map)
+                    return iter(self.weight_map)
 
                 def __len__(self):
-                    return len(weight_map)
+                    return len(self.weight_map)
 
-            state = _State()
+            state = _State(weight_map, shard_paths)
         else:
             path = hf_hub_download(repo, filename, subfolder=subfolder, token=token)
             state = {}
diff --git a/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py b/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py
index de8dbf9e..532fb032 100644
--- a/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py
+++ b/zeromodels/models/qwen_image_21/convert_qwen_image_21_diffusers_to_keras.py
@@ -133,7 +133,6 @@ def load_json(filename, subfolder=None):
     )
 
 
-
 def transfer_qwen_image_21(
     repo, token=None, dtype="bfloat16", build_sample_size=8, config=None
 ):
@@ -198,22 +197,26 @@ def transfer_qwen_image_21(
                 }
 
                 class _State:
+                    def __init__(self, weight_map, shard_paths):
+                        self.weight_map = weight_map
+                        self.shard_paths = shard_paths
+
                     def __contains__(self, key):
-                        return key in weight_map
+                        return key in self.weight_map
 
                     def __getitem__(self, key):
                         with safe_open(
-                            shard_paths[weight_map[key]], framework="np"
+                            self.shard_paths[self.weight_map[key]], framework="np"
                         ) as shard:
                             return shard.get_tensor(key)
 
                     def keys(self):
-                        return weight_map.keys()
+                        return self.weight_map.keys()
 
                     def __iter__(self):
-                        return iter(weight_map)
+                        return iter(self.weight_map)
 
-                state = _State()
+                state = _State(weight_map, shard_paths)
             else:
                 raise FileNotFoundError
         except Exception:
diff --git a/zeromodels/version.py b/zeromodels/version.py
index 1e3de06e..4ec97a1d 100644
--- a/zeromodels/version.py
+++ b/zeromodels/version.py
@@ -1,4 +1,4 @@
-__version__ = "1.3.4"
+__version__ = "1.3.5"
 
 
 def version():