stazizov commited on
Commit
47b72cd
1 Parent(s): 8e44d92

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -26,12 +26,13 @@ else:
26
  @dataclass
27
  class Config:
28
  name: str = "flux-dev"
29
- device: str = "cuda" if torch.cuda.is_available() else "cpu"
30
  offload: bool = False
31
  share: bool = False
32
  ckpt_dir: str = "."
33
 
34
  xflux_pipeline = XFluxPipeline(Config.name, Config.device, Config.offload)
 
35
 
36
  @spaces.GPU
37
  def generate(**kwargs):
 
26
  @dataclass
27
  class Config:
28
  name: str = "flux-dev"
29
+ device: str = "cpu"
30
  offload: bool = False
31
  share: bool = False
32
  ckpt_dir: str = "."
33
 
34
  xflux_pipeline = XFluxPipeline(Config.name, Config.device, Config.offload)
35
+ xflux_pipeline.to(device='cuda' if torch.cuda.is_available() else 'cpu')
36
 
37
  @spaces.GPU
38
  def generate(**kwargs):