modified: src/flux/model.py

src/flux/model.py

@@ -90,6 +90,27 @@ class Flux(nn.Module):
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
 
+        # --- CRITICAL DEBUG: Check the device of self.img_in's parameters ---
+        weight_device = self.img_in.weight.device
+        bias_device = self.img_in.bias.device if self.img_in.bias is not None else "N/A (None)"
+        print(f"self.img_in.weight device: {weight_device}")
+        print(f"self.img_in.bias device: {bias_device}")
+
+        # --- FIX: Explicitly move img to the device of img_in's weight if they differ ---
+        # This is the core fix if the mismatch is here
+        if img.device != weight_device:
+            print(f"!!! Mismatch detected: img on {img.device}, img_in.weight on {weight_device}. Moving img to {weight_device} !!!")
+            img = img.to(weight_device)
+            # It's also good practice to ensure other inputs are on the same device if they aren't already
+            # However, based on your previous check, they should be. Let's double-check one key one:
+            if txt.device != weight_device:
+                    print(f"!!! Also moving 'txt' from {txt.device} to {weight_device} !!!")
+                    txt = txt.to(weight_device)
+            # Add similar checks/moves for txt_ids, y, timesteps, guidance if needed,
+            # but based on your previous debug, they were on cuda:0.
+        
+        print("--- End of Critical Debug ---")
+
         # running on sequences img
         img = self.img_in(img)
         vec = self.time_in(timestep_embedding(timesteps, 256))

src/flux/sampling.py

@@ -522,23 +522,6 @@ def denoise_multi_turn_consistent(
         info['second_order'] = False
         info['inject'] = inject_list[i]
 
-        # Check and print devices of all input tensors
-        tensor_info = {
-            'img': img,
-            'img_ids': img_ids,
-            'txt': txt,
-            'txt_ids': txt_ids,
-            'y (vec)': vec, # Renamed for clarity in print
-            'timesteps (t_vec)': t_vec, # Renamed for clarity in print
-            'guidance (guidance_vec)': guidance_vec # Renamed for clarity in print
-        }
-        for name, tensor in tensor_info.items():
-            if torch.is_tensor(tensor):
-                print(f"{name} device: {tensor.device}, shape: {tensor.shape}")
-            else:
-                print(f"{name} is not a tensor: {type(tensor)} (value: {tensor})")
-
-
         if next_step_velocity is None:
             pred, info = model(
                 img=img,