Initial Commit

app.py

@@ -54,126 +54,126 @@ class FluxInterface:
         self.MODEL_ID = "Djrango/Qwen2vl-Flux"
         
     def load_models(self):
-    if self.models is not None:
-        return
+        if self.models is not None:
+            return
 
-    import gc
-    torch.cuda.empty_cache()
-    gc.collect()
-    
-    logger.info("Starting model loading...")
-    
-    try:
-        # 1. 首先加载小型模型和tokenizer
-        tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer")
-        tokenizer_two = T5TokenizerFast.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer_2")
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.MODEL_ID, subfolder="flux/scheduler", shift=1)
-        
-        # 2. 加载并优化CLIP text encoder
-        text_encoder = CLIPTextModel.from_pretrained(
-            self.MODEL_ID, 
-            subfolder="flux/text_encoder",
-            torch_dtype=self.dtype,
-            device_map="auto"  # 让模型自动管理显存
-        )
-        
-        # 3. 加载T5 encoder
-        text_encoder_two = T5EncoderModel.from_pretrained(
-            self.MODEL_ID,
-            subfolder="flux/text_encoder_2",
-            torch_dtype=self.dtype,
-            device_map="auto"
-        )
-        
-        # 清理一次显存
-        torch.cuda.empty_cache()
-        gc.collect()
-        
-        # 4. 加载VAE
-        vae = AutoencoderKL.from_pretrained(
-            self.MODEL_ID,
-            subfolder="flux/vae",
-            torch_dtype=self.dtype,
-            device_map="auto"
-        )
-        
-        # 5. 加载Transformer
-        transformer = FluxTransformer2DModel.from_pretrained(
-            self.MODEL_ID,
-            subfolder="flux/transformer",
-            torch_dtype=self.dtype,
-            device_map="auto"
-        )
-        
-        # 再次清理显存
+        import gc
         torch.cuda.empty_cache()
         gc.collect()
-        
-        # 6. 加载Qwen2VL
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
-            self.MODEL_ID,
-            subfolder="qwen2-vl",
-            torch_dtype=self.dtype,
-            device_map="auto"
-        )
-        
-        # 7. 加载其他小组件
-        connector = Qwen2Connector().to(self.dtype)
-        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='cpu')
-        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)
+    
+        logger.info("Starting model loading...")
+    
+        try:
+            # 1. 首先加载小型模型和tokenizer
+            tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer")
+            tokenizer_two = T5TokenizerFast.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer_2")
+            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.MODEL_ID, subfolder="flux/scheduler", shift=1)
+            
+            # 2. 加载并优化CLIP text encoder
+            text_encoder = CLIPTextModel.from_pretrained(
+                self.MODEL_ID, 
+                subfolder="flux/text_encoder",
+                torch_dtype=self.dtype,
+                device_map="auto"  # 让模型自动管理显存
+            )
+            
+            # 3. 加载T5 encoder
+            text_encoder_two = T5EncoderModel.from_pretrained(
+                self.MODEL_ID,
+                subfolder="flux/text_encoder_2",
+                torch_dtype=self.dtype,
+                device_map="auto"
+            )
+            
+            # 清理一次显存
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+            # 4. 加载VAE
+            vae = AutoencoderKL.from_pretrained(
+                self.MODEL_ID,
+                subfolder="flux/vae",
+                torch_dtype=self.dtype,
+                device_map="auto"
+            )
+            
+            # 5. 加载Transformer
+            transformer = FluxTransformer2DModel.from_pretrained(
+                self.MODEL_ID,
+                subfolder="flux/transformer",
+                torch_dtype=self.dtype,
+                device_map="auto"
+            )
+            
+            # 再次清理显存
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+            # 6. 加载Qwen2VL
+            qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
+                self.MODEL_ID,
+                subfolder="qwen2-vl",
+                torch_dtype=self.dtype,
+                device_map="auto"
+            )
+            
+            # 7. 加载其他小组件
+            connector = Qwen2Connector().to(self.dtype)
+            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='cpu')
+            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)
 
-        self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype)
-        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='cpu')
-        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)
-        
-        # 设置eval模式和关闭梯度
-        for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]:
-            if hasattr(model, 'eval'):
-                model.eval()
-            if hasattr(model, 'requires_grad_'):
-                model.requires_grad_(False)
-        
-        logger.info("Models loaded successfully")
-        
-        self.models = {
-            'tokenizer': tokenizer,
-            'text_encoder': text_encoder,
-            'text_encoder_two': text_encoder_two,
-            'tokenizer_two': tokenizer_two,
-            'vae': vae,
-            'transformer': transformer,
-            'scheduler': scheduler,
-            'qwen2vl': qwen2vl,
-            'connector': connector
-        }
-        
-        # 初始化processor和pipeline
-        self.qwen2vl_processor = AutoProcessor.from_pretrained(
-            self.MODEL_ID,
-            subfolder="qwen2-vl",
-            min_pixels=256*28*28,
-            max_pixels=256*28*28
-        )
-        
-        self.pipeline = FluxPipeline(
-            transformer=transformer,
-            scheduler=scheduler,
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-        )
-        
-    except Exception as e:
-        logger.error(f"Error loading models: {str(e)}")
-        torch.cuda.empty_cache()
-        gc.collect()
-        raise
+            self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype)
+            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='cpu')
+            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)
+            
+            # 设置eval模式和关闭梯度
+            for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]:
+                if hasattr(model, 'eval'):
+                    model.eval()
+                if hasattr(model, 'requires_grad_'):
+                    model.requires_grad_(False)
+            
+            logger.info("Models loaded successfully")
+            
+            self.models = {
+                'tokenizer': tokenizer,
+                'text_encoder': text_encoder,
+                'text_encoder_two': text_encoder_two,
+                'tokenizer_two': tokenizer_two,
+                'vae': vae,
+                'transformer': transformer,
+                'scheduler': scheduler,
+                'qwen2vl': qwen2vl,
+                'connector': connector
+            }
+            
+            # 初始化processor和pipeline
+            self.qwen2vl_processor = AutoProcessor.from_pretrained(
+                self.MODEL_ID,
+                subfolder="qwen2-vl",
+                min_pixels=256*28*28,
+                max_pixels=256*28*28
+            )
+            
+            self.pipeline = FluxPipeline(
+                transformer=transformer,
+                scheduler=scheduler,
+                vae=vae,
+                text_encoder=text_encoder,
+                tokenizer=tokenizer,
+            )
+            
+        except Exception as e:
+            logger.error(f"Error loading models: {str(e)}")
+            torch.cuda.empty_cache()
+            gc.collect()
+            raise
     
     def resize_image(self, img, max_pixels=1050000):
         if not isinstance(img, Image.Image):