Fix CUBLAS errors: enforce dtype consistency at model entry points

- Add dtype conversion at Flux model forward() entry point
- Add dtype conversion in LastLayer to match weight dtype
- Remove torch.compile (not compatible with ZeroGPU)
- This ensures all inputs match model weights dtype regardless of ZeroGPU auto-conversion

app.py

@@ -148,6 +148,7 @@ def init_generator():
             use_4bit_quantization=False,  # Full precision model
             dtype="fp32",  # Use fp32 to avoid CUBLAS errors on ZeroGPU
         )
+            
     return generator
 
 

src/flux/model.py

@@ -199,6 +199,30 @@ class Flux(nn.Module):
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
 
+        # Get model's weight dtype and ensure all inputs match
+        # This prevents CUBLAS errors on ZeroGPU where tensors may be auto-converted
+        weight_dtype = self.img_in.weight.dtype
+        if img.dtype != weight_dtype:
+            img = img.to(weight_dtype)
+        if img_ids.dtype != weight_dtype:
+            img_ids = img_ids.to(weight_dtype)
+        if txt.dtype != weight_dtype:
+            txt = txt.to(weight_dtype)
+        if txt_ids.dtype != weight_dtype:
+            txt_ids = txt_ids.to(weight_dtype)
+        if y.dtype != weight_dtype:
+            y = y.to(weight_dtype)
+        if timesteps.dtype != weight_dtype:
+            timesteps = timesteps.to(weight_dtype)
+        if timesteps2 is not None and timesteps2.dtype != weight_dtype:
+            timesteps2 = timesteps2.to(weight_dtype)
+        if cond_txt_latent is not None and cond_txt_latent.dtype != weight_dtype:
+            cond_txt_latent = cond_txt_latent.to(weight_dtype)
+        if guidance is not None and guidance.dtype != weight_dtype:
+            guidance = guidance.to(weight_dtype)
+        if image_proj is not None and image_proj.dtype != weight_dtype:
+            image_proj = image_proj.to(weight_dtype)
+
         # running on sequences img
         img = self.img_in(img)
         if self.module_embeddings is not None:

src/flux/modules/layers.py

@@ -625,6 +625,17 @@ class LastLayer(nn.Module):
         self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
 
     def forward(self, x: Tensor, vec: Tensor, vec2: Tensor) -> Tensor:
+        # Get weights dtype (should be fp32 if loaded correctly)
+        weight_dtype = self.linear.weight.dtype
+        
+        # Ensure inputs match weight dtype to avoid CUBLAS errors on ZeroGPU
+        if x.dtype != weight_dtype:
+            x = x.to(weight_dtype)
+        if vec.dtype != weight_dtype:
+            vec = vec.to(weight_dtype)
+        if vec2 is not None and vec2.dtype != weight_dtype:
+            vec2 = vec2.to(weight_dtype)
+        
         shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
         shift = shift[:, None, :]
         scale = scale[:, None, :]