Update generator.py

generator.py

@@ -751,18 +751,18 @@ class RetroArtConverter:
         
         pipe_kwargs["generator"] = generator
         
-        # --- START FIX: Handle CPU Offload for Compel ---
         # Use Compel for prompt encoding if available
         if self.use_compel and self.compel is not None:
             try:
                 print("Encoding prompts with Compel...")
                 
-                # Move text encoders to GPU
+                # --- FIX 1: Move encoders to GPU ---
                 self.pipe.text_encoder.to(self.device)
                 self.pipe.text_encoder_2.to(self.device)
                 
-                conditioning = self.compel(prompt)
-                negative_conditioning = self.compel(negative_prompt)
+                # --- FIX 2: Tell Compel to use the GPU for its tensors ---
+                conditioning = self.compel(prompt, device=self.device)
+                negative_conditioning = self.compel(negative_prompt, device=self.device)
                 
                 pipe_kwargs["prompt_embeds"] = conditioning[0]
                 pipe_kwargs["pooled_prompt_embeds"] = conditioning[1]
@@ -775,7 +775,7 @@ class RetroArtConverter:
                 pipe_kwargs["prompt"] = prompt
                 pipe_kwargs["negative_prompt"] = negative_prompt
             finally:
-                # Move text encoders back to CPU
+                # --- FIX 3: Move encoders back to CPU ---
                 try:
                     self.pipe.text_encoder.to("cpu")
                     self.pipe.text_encoder_2.to("cpu")
@@ -784,7 +784,6 @@ class RetroArtConverter:
         else:
             pipe_kwargs["prompt"] = prompt
             pipe_kwargs["negative_prompt"] = negative_prompt
-        # --- END FIX ---
         
         # Add CLIP skip
         if hasattr(self.pipe, 'text_encoder'):