Update pipeline_qwenimage_edit.py

pipeline_qwenimage_edit.py

@@ -170,6 +170,29 @@ def resize_to_multiple_of(image, multiple_of=32):
     return image
 
 
+def ensure_device_recursive(model, target_device):
+    """
+    递归确保模型的所有组件（包括 buffer）都在目标设备上。
+    这是为了解决 ZeroGPU 环境下 register_buffer 的张量可能不会被正确移动的问题。
+    """
+    if target_device is None:
+        return
+    
+    target_device = torch.device(target_device)
+    
+    for name, module in model.named_modules():
+        # 移动所有 buffer
+        for buf_name, buf in list(module._buffers.items()):
+            if buf is not None and buf.device != target_device:
+                module._buffers[buf_name] = buf.to(target_device)
+        
+        # 特别处理 RoPE 相关的属性（有些可能不是通过 register_buffer 注册的）
+        for attr_name in ['inv_freq', 'cos_cached', 'sin_cached', '_cos_cached', '_sin_cached']:
+            if hasattr(module, attr_name):
+                attr = getattr(module, attr_name)
+                if isinstance(attr, torch.Tensor) and attr.device != target_device:
+                    setattr(module, attr_name, attr.to(target_device))
+
 
 class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
     r"""
@@ -223,6 +246,43 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
         self.prompt_template_encode_start_idx = 64
         self.default_sample_size = 128
+        
+        # ZeroGPU 兼容性标记
+        self._device_ensured = False
+
+    def _ensure_device_consistency(self):
+        """
+        确保所有模型组件在同一设备上。
+        在 ZeroGPU 环境下，这个方法应该在 @spaces.GPU 装饰的函数内部调用。
+        """
+        if self._device_ensured:
+            return
+            
+        # 获取目标设备
+        try:
+            target_device = next(self.text_encoder.parameters()).device
+        except StopIteration:
+            return
+        
+        if target_device.type != 'cuda':
+            return
+            
+        print(f"   [PIPELINE] Ensuring device consistency on {target_device}...")
+        
+        # 确保所有模型的 buffer 都在正确的设备上
+        ensure_device_recursive(self.text_encoder, target_device)
+        ensure_device_recursive(self.transformer, target_device)
+        ensure_device_recursive(self.vae, target_device)
+        
+        self._device_ensured = True
+        print(f"   [PIPELINE] Device consistency ensured.")
+
+    def to(self, *args, **kwargs):
+        """
+        重写 to 方法，重置设备一致性标记
+        """
+        self._device_ensured = False
+        return super().to(*args, **kwargs)
 
     # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
@@ -240,10 +300,21 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         device: Optional[torch.device] = None,
         dtype: Optional[torch.dtype] = None,
     ):
-        # 修复：优先使用 text_encoder 的实际设备，而不是 self.device
-        # 这确保在 ZeroGPU 环境下设备一致性
+        # 确保设备一致性（ZeroGPU 修复）
+        self._ensure_device_consistency()
+        
+        # 获取 text_encoder 的实际设备
+        encoder_device = next(self.text_encoder.parameters()).device
         if device is None:
-            device = next(self.text_encoder.parameters()).device
+            device = encoder_device
+        else:
+            # 确保 device 是 torch.device 对象
+            device = torch.device(device)
+            # 如果指定的设备与 encoder 设备不同，使用 encoder 的设备
+            if device != encoder_device:
+                print(f"   [WARNING] Requested device {device} differs from encoder device {encoder_device}, using encoder device")
+                device = encoder_device
+                
         dtype = dtype or self.text_encoder.dtype
 
         prompts = [prompts] if isinstance(prompts, str) else prompts
@@ -285,16 +356,19 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             return_tensors="pt"
         )
         
-        # 修复：明确将每个张量移动到正确的设备
-        # 不依赖 .to(device) 的自动传播
+        # 关键修复：确保所有输入张量都在正确的设备上
         input_ids = model_inputs.input_ids.to(device)
         attention_mask = model_inputs.attention_mask.to(device)
         pixel_values = model_inputs.pixel_values.to(device=device, dtype=dtype)
         image_grid_thw = model_inputs.image_grid_thw.to(device)
 
-        # template = self.prompt_template_encode
+        # 调试输出
+        print(f"   [DEBUG] Input devices - input_ids: {input_ids.device}, pixel_values: {pixel_values.device}")
+        print(f"   [DEBUG] Encoder device: {encoder_device}")
+
         drop_idx = self.prompt_template_encode_start_idx
 
+        # 调用 text_encoder
         outputs = self.text_encoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -343,7 +417,7 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                 provided, text embeddings will be generated from `prompt` input argument.
         """
-        # 修复：优先使用 text_encoder 的实际设备
+        # 获取 text_encoder 的实际设备
         if device is None:
             device = next(self.text_encoder.parameters()).device
 
@@ -671,6 +745,9 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
+        # ZeroGPU 修复：确保设备一致性
+        self._ensure_device_consistency()
+        
         if not isinstance(images, (list, tuple)):
             images = [images]
 
@@ -714,7 +791,7 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         else:
             batch_size = prompt_embeds.shape[0]
 
-        # 修复：使用 text_encoder 的实际设备，而不是 _execution_device
+        # 获取 text_encoder 的实际设备
         device = next(self.text_encoder.parameters()).device
         
         # 3. Preprocess image