Spaces:

MisterAI
/

GenDoc_05

Sleeping

MisterAI commited on Feb 15

Commit

9d225f0

verified ·

1 Parent(s): 6e63c31

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -80,7 +80,7 @@ class FluxGenerator:
             self.pipeline = FluxPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-schnell",
                 revision="refs/pr/1",
-                device_map="auto",
                 torch_dtype=torch.float32,
                 use_safetensors=True
 )

             self.pipeline = FluxPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-schnell",
                 revision="refs/pr/1",
+                device_map="balanced",
                 torch_dtype=torch.float32,
                 use_safetensors=True
 )