32 lines
1017 B
Python
32 lines
1017 B
Python
import torch
|
|
from diffusers import DiffusionPipeline
|
|
|
|
# Load a small test to see how pipeline handles latents
|
|
print("Testing latents device handling...")
|
|
|
|
# Simulate what we're doing
|
|
class MockPipeline:
|
|
def __init__(self):
|
|
self.device = torch.device("cuda")
|
|
self._execution_device = torch.device("cuda")
|
|
|
|
# Our current approach
|
|
pipeline_mock = MockPipeline()
|
|
|
|
# Get device (our current code)
|
|
if hasattr(pipeline_mock, "_execution_device"):
|
|
device = pipeline_mock._execution_device
|
|
else:
|
|
device = pipeline_mock.device
|
|
|
|
print(f"Device selected: {device}")
|
|
|
|
# Create latents (simulating our VAE encode result)
|
|
latents = torch.randn(1, 4, 128, 128, device=device, dtype=torch.bfloat16)
|
|
print(f"Latents device: {latents.device}")
|
|
print(f"Latents dtype: {latents.dtype}")
|
|
|
|
# When we pass to pipeline, it should already be on the correct device
|
|
print(f"\n✓ Latents are already on the execution device: {device}")
|
|
print("✓ Pipeline will accept them directly without device transfers")
|