perf

app.py

@@ -16,18 +16,20 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 #
 # login()
 
-dtype = torch.bfloat16
+dtype = torch.float16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 base_model = "black-forest-labs/FLUX.1-dev"
 controlnet_model = "YishaoAI/flux-dev-controlnet-canny-kid-clothes"
 
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=dtype)
+controlnet = FluxControlNetModel.from_pretrained(
+    controlnet_model, torch_dtype=dtype, device_map="auto"
+)
 pipe = FluxControlNetInpaintPipeline.from_pretrained(
-    base_model, controlnet=controlnet, torch_dtype=dtype
+    base_model, controlnet=controlnet, torch_dtype=dtype, device_map="auto"
 ).to(device)
 
-# pipe.enable_model_cpu_offload()
+pipe.enable_model_cpu_offload()
 
 canny = CannyDetector()