package

requirements.txt

@@ -1,19 +1,17 @@
---extra-index-url https://download.pytorch.org/whl/cu124
 torch==2.1.0
+torchdiffeq==0.2.5
 torchvision==0.16.0
-numpy<2
+numpy==1.26.3
 diffusers==0.32.1
-accelerate==1.1.1
-transformers==4.46.2
+accelerate==1.2.1
+transformers==4.47.1
 huggingface-hub==0.25.0 
 tensorboard
 gradio
-torchdiffeq
 click
-torchvision
 opencv-python
 scikit-image
-numba
+numba==0.60.0
 scipy
 tqdm
 einops

visualcloze.py

@@ -12,6 +12,7 @@ from flux.util import load_clip, load_t5, load_flow_model
 from transport import Sampler, create_transport
 from imgproc import to_rgb_if_rgba
 
+
 def center_crop(image, target_size):
     width, height = image.size
     new_width, new_height = target_size
@@ -90,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(