deepseekmirror/Janus
1
0
Fork 0
mirror of https://github.com/deepseek-ai/Janus.git synced 2025-02-22 13:48:57 -05:00
Code Issues Actions Packages Projects Releases Wiki Activity

更新CUDA设备管理逻辑，支持MPS设备

Browse Source
This commit is contained in:
censujiang 2025-01-31 04:01:37 +08:00
parent 581fdd1489
commit eb441468d8
3 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
demo/app.py
View File

@ -19,7 +19,7 @@ vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
cuda_device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
# Multimodal Understanding function
@torch.inference_mode()
# Multimodal Understanding function

8
demo/app_janusflow.py
View File

@ -5,7 +5,13 @@ from PIL import Image
from diffusers.models import AutoencoderKL
import numpy as np
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
cuda_device = 'cpu'
if torch.cuda.is_available():
cuda_device = 'cuda'
elif torch.backends.mps.is_available():
cuda_device = 'mps'
else:
cuda_device = 'cpu'
# Load model and processor
model_path = "deepseek-ai/JanusFlow-1.3B"

2
demo/fastapi_app.py
View File

@ -21,7 +21,7 @@ vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
cuda_device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
@torch.inference_mode()