Hugging Face's logo

NoMoreCopyrightOrg
/
flux-dev

Text-to-Image
Diffusers
Safetensors
English
FluxPipeline
Model card Files
xet
 Community
John6666 commited on
Commit
31523d1
·
verified ·
1 Parent(s): acf2d5f

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -32,7 +32,7 @@ class EndpointHandler:
32
  self.pipeline.transformer.fuse_qkv_projections()
33
  self.pipeline.vae.fuse_qkv_projections()
34
  if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
35
- self.pipeline.to(distributed_state.device)
36
 
37
  @torch.inference_mode()
38
  def __call__(self, data: Dict[str, Any]) -> Image.Image:
 
32
  self.pipeline.transformer.fuse_qkv_projections()
33
  self.pipeline.vae.fuse_qkv_projections()
34
  if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
35
+ self.pipeline.to("cuda")
36
 
37
  @torch.inference_mode()
38
  def __call__(self, data: Dict[str, Any]) -> Image.Image: