mrfakename commited on
Commit
d20ce72
·
verified ·
1 Parent(s): 4111425

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -11,7 +11,7 @@ MODEL_CACHE = {}
11
  for id, model in MODELS.items():
12
  print(f"Loading model {model}...")
13
  MODEL_CACHE[model] = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
14
- pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
15
  print(f"Loaded model {model}")
16
 
17
  @spaces.GPU
 
11
  for id, model in MODELS.items():
12
  print(f"Loading model {model}...")
13
  MODEL_CACHE[model] = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
14
+ MODEL_CACHE[model].enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
15
  print(f"Loaded model {model}")
16
 
17
  @spaces.GPU