Spaces:
Build error
Build error
Update app.py
Browse files
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 |
-
|
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
|