space

visualcloze.py

@@ -91,26 +91,26 @@ class VisualClozeModel:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[self.precision]
         
-        # # Initialize model
-        # print("Initializing model...")
-        # self.model = load_flow_model(model_name, device=self.device, lora_rank=self.lora_rank)
+        # Initialize model
+        print("Initializing model...")
+        self.model = load_flow_model(model_name, device=self.device, lora_rank=self.lora_rank)
         
-        # # Initialize VAE
-        # print("Initializing VAE...")
-        # self.ae = AutoencoderKL.from_pretrained(f"black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=self.dtype).to(self.device)
-        # self.ae.requires_grad_(False)
+        # Initialize VAE
+        print("Initializing VAE...")
+        self.ae = AutoencoderKL.from_pretrained(f"black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=self.dtype).to(self.device)
+        self.ae.requires_grad_(False)
         
-        # # Initialize text encoders
-        # print("Initializing text encoders...")
-        # self.t5 = load_t5(self.device, max_length=self.max_length)
-        # self.clip = load_clip(self.device)
+        # Initialize text encoders
+        print("Initializing text encoders...")
+        self.t5 = load_t5(self.device, max_length=self.max_length)
+        self.clip = load_clip(self.device)
         
-        # self.model.eval().to(self.device, dtype=self.dtype)
+        self.model.eval().to(self.device, dtype=self.dtype)
         
-        # # Load model weights
-        # ckpt = torch.load(model_path)
-        # self.model.load_state_dict(ckpt, strict=False)
-        # del ckpt
+        # Load model weights
+        ckpt = torch.load(model_path)
+        self.model.load_state_dict(ckpt, strict=False)
+        del ckpt
         
         # Initialize sampler
         transport = create_transport(