NoMoreCopyrightOrg
/

flux-dev

Inference Endpoints

Model card Files Files and versions Community

John6666 commited on 6 days ago

Commit

31523d1

·

verified ·

1 Parent(s): acf2d5f

Upload handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -32,7 +32,7 @@ class EndpointHandler:
         self.pipeline.transformer.fuse_qkv_projections()
         self.pipeline.vae.fuse_qkv_projections()
         if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
-        self.pipeline.to(distributed_state.device)
     @torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> Image.Image:

         self.pipeline.transformer.fuse_qkv_projections()
         self.pipeline.vae.fuse_qkv_projections()
         if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
+        self.pipeline.to("cuda")
     @torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> Image.Image: