FluxM-Lightning-Upscaler

Running on Zero

LPX commited on Jun 12

Commit

ec75dde

1 Parent(s): c80eda9

Add GPU decorator to safe_model_load function in model_loader.py

Files changed (1) hide show

model_loader.py CHANGED Viewed

@@ -1,10 +1,12 @@
 # model_loader.py
 import os
 import torch
 from diffusers import FluxControlNetPipeline
 from transformers import T5EncoderModel
 from moondream import vl
 def safe_model_load():
     """Load models in a single GPU invocation to keep them warm"""
     try:

 # model_loader.py
 import os
 import torch
+import spaces
 from diffusers import FluxControlNetPipeline
 from transformers import T5EncoderModel
 from moondream import vl
+@spaces.GPU()
 def safe_model_load():
     """Load models in a single GPU invocation to keep them warm"""
     try: