Initial Commit

app.py

@@ -11,6 +11,9 @@ import math
 import logging
 import sys
 
+import os
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512,expandable_segments:True'
+
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 
 # 设置日志
@@ -51,49 +54,92 @@ 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
 
-        logger.info("Starting model loading...")
-        
-        # Load FLUX components
+    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")
-        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
-        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
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(self.MODEL_ID, subfolder="qwen2-vl").to(self.dtype).to(self.device)
+        # 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"
+        )
         
-        # Load connector
-        connector = Qwen2Connector().to(self.dtype).to(self.device)
+        # 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')
-        # 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).to(self.dtype).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')
-        # 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)
         
-        # Set models to eval mode
+        # 设置eval模式和关闭梯度
         for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]:
-            model.requires_grad_(False)
-            model.eval()
+            if hasattr(model, 'eval'):
+                model.eval()
+            if hasattr(model, 'requires_grad_'):
+                model.requires_grad_(False)
         
-        logger.info("All models loaded successfully")
+        logger.info("Models loaded successfully")
         
         self.models = {
             'tokenizer': tokenizer,
@@ -107,11 +153,11 @@ class FluxInterface:
             'connector': connector
         }
         
-        # Initialize processor and pipeline
+        # 初始化processor和pipeline
         self.qwen2vl_processor = AutoProcessor.from_pretrained(
-            self.MODEL_ID, 
+            self.MODEL_ID,
             subfolder="qwen2-vl",
-            min_pixels=256*28*28, 
+            min_pixels=256*28*28,
             max_pixels=256*28*28
         )
         
@@ -121,7 +167,13 @@ class FluxInterface:
             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):