Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2025-12-25 12:39:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

Fix device_map for DiffusionForTextToImageSynthesis

Browse Source 
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13115051
* add kwargs for device map
This commit is contained in:
hemu.zp
2023-06-29 16:49:20 +08:00
committed by wenmeng.zwm
parent d4b9396084
commit 1944416d3d
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
modelscope/models/multi_modal/diffusion/model.py
View File

@@ -114,9 +114,9 @@ class DiffusionModel(nn.Module):
Tasks.text_to_image_synthesis, module_name=Models.diffusion)
class DiffusionForTextToImageSynthesis(Model):
def __init__(self, model_dir, device='gpu'):
def __init__(self, model_dir, device='gpu', **kwargs):
device = 'gpu' if torch.cuda.is_available() else 'cpu'
super().__init__(model_dir=model_dir, device=device)
super().__init__(model_dir=model_dir, device=device, **kwargs)
diffusion_model = DiffusionModel(model_dir=model_dir)
pretrained_params = torch.load(
osp.join(model_dir, ModelFile.TORCH_MODEL_BIN_FILE), 'cpu')