Update app.py

app.py

@@ -20,7 +20,7 @@ class ModelWrapper:
         super().__init__()
         torch.set_grad_enabled(False)
         
-        self.DTYPE = getattr(torch, precision)
+        self.DTYPE = torch.float16
         self.device = accelerator.device
 
         self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
@@ -49,7 +49,7 @@ class ModelWrapper:
     
     def create_generator(self, model_id, checkpoint_path):
         generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
-        state_dict = torch.load(checkpoint_path, map_location="cuda")
+        state_dict = torch.load(checkpoint_path)
         generator.load_state_dict(state_dict, strict=True)
         generator.requires_grad_(False)
         return generator