ACE-Plus

Running on Zero

chaojiemao commited on Dec 8, 2024

Commit

532ddca

verified ·

1 Parent(s): 3c39b17

Update model/flux.py

Files changed (1) hide show

model/flux.py CHANGED Viewed

@@ -652,7 +652,7 @@ class Flux(BaseModel):
     def load_pretrained_model(self, pretrained_model):
         if next(self.parameters()).device.type == 'meta':
-            map_location = we.device_id
         else:
             map_location = "cpu"
         if self.lora_model is not None:

     def load_pretrained_model(self, pretrained_model):
         if next(self.parameters()).device.type == 'meta':
+            map_location = torch.device(we.device_id)
         else:
             map_location = "cpu"
         if self.lora_model is not None: