Align with pulid demo

app.py

@@ -431,6 +431,7 @@ if __name__ == "__main__":
         args.offload = True
 
     print(f"Using device: {args.device}")
+    print(f"fp8: {args.fp8}")
     print(f"Offload: {args.offload}")
 
     demo = create_demo(args, args.name, args.device, args.offload, args.aggressive_offload)

flux/util.py

@@ -123,36 +123,17 @@ def load_flow_model(name: str, device: str = "cuda", hf_download: bool = True):
     ):
         ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow, local_dir='models')
 
-# Initialize the model on the 'meta' device, which doesn't allocate real memory
-    with torch.device('meta'):
-        model = Flux(configs[name].params)
-    model = model.to_empty(device=device)
+    with torch.device(device):
+        model = Flux(configs[name].params).to(torch.bfloat16)
 
     if ckpt_path is not None:
         print("Loading checkpoint")
-        # Load the state dictionary directly to the desired device
+        # load_sft doesn't support torch.device
         sd = load_sft(ckpt_path, device=str(device))
-        # Load the state dictionary into the model
         missing, unexpected = model.load_state_dict(sd, strict=False)
         print_load_warning(missing, unexpected)
-    model.to(torch.bfloat16)
     return model
 
-# from XLabs-AI https://github.com/XLabs-AI/x-flux/blob/1f8ef54972105ad9062be69fe6b7f841bce02a08/src/flux/util.py#L330
-def load_flow_model_quintized(name: str, device: str = "cuda", hf_download: bool = True):
-    # Loading Flux
-    print("Init model")
-    ckpt_path = 'models/flux-dev-fp8.safetensors'
-    if (
-        not os.path.exists(ckpt_path)
-        and hf_download
-    ):
-        print("Downloading model")
-        ckpt_path = hf_hub_download("XLabs-AI/flux-dev-fp8", "flux-dev-fp8.safetensors")
-        print("Model downloaded to", ckpt_path)
-    json_path = hf_hub_download("XLabs-AI/flux-dev-fp8", 'flux_dev_quantization_map.json')
-
-    model = Flux(configs[name].params).to(torch.bfloat16)
 def load_flow_model_quintized(
     name: str,
     device: str = "cuda",