Spaces:
Running
on
Zero
Running
on
Zero
File size: 941 Bytes
f876753 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
from typing import *
import torch
import torch.nn as nn
from diffusers import DiffusionPipeline
class Pipeline(DiffusionPipeline):
"""Base class for all pipelines."""
def __init__(self):
super().__init__()
def __call__(self, *args, **kwargs):
raise NotImplementedError
def enable_xformers_memory_efficient_attention(self):
pass
def enable_model_cpu_offload(self):
pass
@property
def device(self) -> torch.device:
for model in self.models.values():
if hasattr(model, 'device'):
return model.device
for model in self.models.values():
if hasattr(model, 'parameters'):
return next(model.parameters()).device
raise RuntimeError("No device found.")
def to(self, device: torch.device) -> None:
for model in self.models.values():
model.to(device) |