mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
83 lines
3.9 KiB
Python
83 lines
3.9 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from enum import Enum
|
|
from invokeai.backend.model_management.model_cache import ModelCache
|
|
|
|
class DummyModelBase(object):
|
|
'''Base class for dummy component of a diffusers model'''
|
|
def __init__(self, repo_id):
|
|
self.repo_id = repo_id
|
|
self.device = torch.device('cpu')
|
|
|
|
@classmethod
|
|
def from_pretrained(cls,
|
|
repo_id:str,
|
|
revision:str=None,
|
|
subfolder:str=None,
|
|
cache_dir:str=None,
|
|
):
|
|
return cls(repo_id)
|
|
|
|
def to(self, device):
|
|
self.device = device
|
|
|
|
class DummyModelType1(DummyModelBase):
|
|
pass
|
|
|
|
class DummyModelType2(DummyModelBase):
|
|
pass
|
|
|
|
class DummyPipeline(DummyModelBase):
|
|
'''Dummy pipeline object is a composite of several types'''
|
|
def __init__(self,repo_id):
|
|
super().__init__(repo_id)
|
|
self.type1 = DummyModelType1('dummy/type1')
|
|
self.type2 = DummyModelType2('dummy/type2')
|
|
|
|
class DMType(Enum):
|
|
dummy_pipeline = DummyPipeline
|
|
type1 = DummyModelType1
|
|
type2 = DummyModelType2
|
|
|
|
cache = ModelCache(max_cache_size=4)
|
|
|
|
def test_pipeline_fetch():
|
|
assert cache.cache_size()==0
|
|
with cache.get_model('dummy/pipeline1',DMType.dummy_pipeline) as pipeline1,\
|
|
cache.get_model('dummy/pipeline1',DMType.dummy_pipeline) as pipeline1a,\
|
|
cache.get_model('dummy/pipeline2',DMType.dummy_pipeline) as pipeline2:
|
|
assert pipeline1 is not None, 'get_model() should not return None'
|
|
assert pipeline1a is not None, 'get_model() should not return None'
|
|
assert pipeline2 is not None, 'get_model() should not return None'
|
|
assert type(pipeline1)==DMType.dummy_pipeline.value,'get_model() did not return model of expected type'
|
|
assert pipeline1==pipeline1a,'pipelines with the same repo_id should be the same'
|
|
assert pipeline1!=pipeline2,'pipelines with different repo_ids should not be the same'
|
|
assert len(cache.models)==2,'cache should uniquely cache models with same identity'
|
|
with cache.get_model('dummy/pipeline3',DMType.dummy_pipeline) as pipeline3,\
|
|
cache.get_model('dummy/pipeline4',DMType.dummy_pipeline) as pipeline4:
|
|
assert len(cache.models)==4,'cache did not grow as expected'
|
|
|
|
def test_signatures():
|
|
with cache.get_model('dummy/pipeline',DMType.dummy_pipeline,revision='main') as pipeline1,\
|
|
cache.get_model('dummy/pipeline',DMType.dummy_pipeline,revision='fp16') as pipeline2,\
|
|
cache.get_model('dummy/pipeline',DMType.dummy_pipeline,revision='main',subfolder='foo') as pipeline3:
|
|
assert pipeline1 != pipeline2,'models are distinguished by their revision'
|
|
assert pipeline1 != pipeline3,'models are distinguished by their subfolder'
|
|
|
|
def test_pipeline_device():
|
|
with cache.get_model('dummy/pipeline1',DMType.type1) as model1:
|
|
assert model1.device==torch.device('cuda'),'when in context, model device should be in GPU'
|
|
with cache.get_model('dummy/pipeline1',DMType.type1, gpu_load=False) as model1:
|
|
assert model1.device==torch.device('cpu'),'when gpu_load=False, model device should be CPU'
|
|
|
|
def test_submodel_fetch():
|
|
with cache.get_model(repo_id_or_path='dummy/pipeline1',model_type=DMType.dummy_pipeline) as pipeline,\
|
|
cache.get_model(repo_id_or_path='dummy/pipeline1',model_type=DMType.dummy_pipeline,submodel=DMType.type1) as part1,\
|
|
cache.get_model(repo_id_or_path='dummy/pipeline2',model_type=DMType.dummy_pipeline,submodel=DMType.type1) as part2:
|
|
assert type(part1)==DummyModelType1,'returned submodel is not of expected type'
|
|
assert part1.device==torch.device('cuda'),'returned submodel should be in the GPU when in context'
|
|
assert pipeline.type1==part1,'returned submodel should match the corresponding subpart of parent model'
|
|
assert pipeline.type1!=part2,'returned submodel should not match the subpart of a different parent'
|
|
|