Update modelmod.py

modelmod.py

@@ -18,7 +18,7 @@ import os
 import cv2
 import numpy as np
 import math
-nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
+nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 
@@ -145,8 +145,8 @@ class FluxModel:
         self.transformer = FluxTransformer2DModel.from_pretrained(MODEL_PATHS['flux'], subfolder="transformer", token=HF_TOKEN, **self.qkwargs)
 
         self.vae.requires_grad_(False).to(self.dtype).to(self.device)
-        self.transformer.requires_grad_(False).to(self.dtype).to(self.device)
-        #self.transformer.requires_grad_(False).to(self.device)
+        #self.transformer.requires_grad_(False).to(self.dtype).to(self.device)
+        self.transformer.requires_grad_(False).to(self.device)
 
     def _init_controlnet(self):
         """Initialize ControlNet model"""