Skip to content

Commit

Permalink
Refactor device variable usage and specify imageio[ffmpeg] version
Browse files Browse the repository at this point in the history
  • Loading branch information
FlyingFathead committed Oct 13, 2024
1 parent 2397c78 commit 45749bf
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions scripts/demo/streamlit_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,15 +67,15 @@ def init_st(version_dict, load_ckpt=True, load_filter=True):
return state

def load_model(model):
device = torch.device("cuda" if USE_CUDA and torch.cuda.is_available() else "cpu")
global device # Use the global device variable
model.to(device)

def set_lowvram_mode(mode):
global lowvram_mode
lowvram_mode = mode

def initial_model_load(model):
device = torch.device("cuda" if USE_CUDA and torch.cuda.is_available() else "cpu")
global device # Use the global device variable
global lowvram_mode
if lowvram_mode:
model.model.half().to(device)
Expand Down

0 comments on commit 45749bf

Please sign in to comment.