Initial Commit

app.py

@@ -52,38 +52,47 @@ class FluxInterface:
         if self.models is not None:
             return
 
+        logger.info("Starting model loading...")
+        
         # Load FLUX components
         tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer")
-        text_encoder = CLIPTextModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder")
-        text_encoder_two = T5EncoderModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder_2")
+        text_encoder = CLIPTextModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder").to(self.dtype).to(self.device)
+        text_encoder_two = T5EncoderModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder_2").to(self.dtype).to(self.device)
         tokenizer_two = T5TokenizerFast.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer_2")
         
-        # Load VAE and transformer from flux folder
-        vae = AutoencoderKL.from_pretrained(self.MODEL_ID, subfolder="flux/vae")
-        transformer = FluxTransformer2DModel.from_pretrained(self.MODEL_ID, subfolder="flux/transformer")
+        # Load VAE and transformer
+        vae = AutoencoderKL.from_pretrained(self.MODEL_ID, subfolder="flux/vae").to(self.dtype).to(self.device)
+        transformer = FluxTransformer2DModel.from_pretrained(self.MODEL_ID, subfolder="flux/transformer").to(self.dtype).to(self.device)
         scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.MODEL_ID, subfolder="flux/scheduler", shift=1)
         
-        # Load Qwen2VL components from qwen2-vl folder
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(self.MODEL_ID, subfolder="qwen2-vl")
+        # Load Qwen2VL components
+        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(self.MODEL_ID, subfolder="qwen2-vl").to(self.dtype).to(self.device)
         
-        # Load connector and t5 embedder from qwen2-vl folder
-        connector = Qwen2Connector()
+        # Load connector
+        connector = Qwen2Connector().to(self.dtype).to(self.device)
         connector_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/connector.pt"
-        connector_state = torch.hub.load_state_dict_from_url(connector_path, map_location=self.device)
+        connector_state = torch.hub.load_state_dict_from_url(connector_path, map_location='cpu')
+        # Move state dict to dtype before loading
+        connector_state = {k: v.to(self.dtype) for k, v in connector_state.items()}
         connector.load_state_dict(connector_state)
+        connector = connector.to(self.device)
         
         # Load T5 embedder
-        self.t5_context_embedder = nn.Linear(4096, 3072)
+        self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype).to(self.device)
         t5_embedder_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/t5_embedder.pt"
-        t5_embedder_state = torch.hub.load_state_dict_from_url(t5_embedder_path, map_location=self.device)
+        t5_embedder_state = torch.hub.load_state_dict_from_url(t5_embedder_path, map_location='cpu')
+        # Move state dict to dtype before loading
+        t5_embedder_state = {k: v.to(self.dtype) for k, v in t5_embedder_state.items()}
         self.t5_context_embedder.load_state_dict(t5_embedder_state)
+        self.t5_context_embedder = self.t5_context_embedder.to(self.device)
         
-        # Move models to device and set dtype
-        models = [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]
-        for model in models:
-            model.to(self.device).to(self.dtype)
+        # Set models to eval mode
+        for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]:
+            model.requires_grad_(False)
             model.eval()
         
+        logger.info("All models loaded successfully")
+        
         self.models = {
             'tokenizer': tokenizer,
             'text_encoder': text_encoder,
@@ -110,7 +119,7 @@ class FluxInterface:
             vae=vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-        )
+        ) 
     
     def resize_image(self, img, max_pixels=1050000):
         if not isinstance(img, Image.Image):