Fix: diffusers from git main for Flux2KleinPipeline, low_cpu_mem_usage=False, no cpu_offload

sprite_lora_resume_v6.py

@@ -2,7 +2,7 @@
 # requires-python = ">=3.10"
 # dependencies = [
 #   "torch",
-#   "diffusers>=0.32.0",
+#   "diffusers @ git+https://github.com/huggingface/diffusers.git",
 #   "transformers",
 #   "accelerate",
 #   "peft",
@@ -211,10 +211,7 @@ def collate_fn(batch):
 
 def train(token):
     """Main training loop."""
-    try:
-        from diffusers import Flux2KleinPipeline as KleinPipeline
-    except ImportError:
-        from diffusers import Flux2Pipeline as KleinPipeline
+    from diffusers import Flux2KleinPipeline
     from peft import LoraConfig, get_peft_model
     from accelerate import Accelerator
 
@@ -234,18 +231,17 @@ def train(token):
 
     # --- Load model ---
     accelerator.print(f"Loading {BASE_MODEL}...")
-    pipe = KleinPipeline.from_pretrained(
+    pipe = Flux2KleinPipeline.from_pretrained(
         BASE_MODEL,
         torch_dtype=torch.bfloat16,
         low_cpu_mem_usage=False,  # Materialize all weights (prevents meta tensor crash)
     )
     # Do NOT call pipe.enable_model_cpu_offload() — accelerator manages device placement.
-    # Move frozen components (VAE, text encoders) to GPU manually.
+    # Move frozen components to GPU manually.
+    # FLUX.2-klein uses Qwen3 text encoder + AutoencoderKLFlux2
     pipe.vae = pipe.vae.to(device)
     if hasattr(pipe, 'text_encoder') and pipe.text_encoder is not None:
         pipe.text_encoder = pipe.text_encoder.to(device)
-    if hasattr(pipe, 'text_encoder_2') and pipe.text_encoder_2 is not None:
-        pipe.text_encoder_2 = pipe.text_encoder_2.to(device)
 
     # --- Apply LoRA ---
     lora_config = LoraConfig(