Initial Commit

app.py

@@ -12,7 +12,9 @@ import logging
 import sys
 
 import os
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512,expandable_segments:True'
+# 设置环境变量，强制禁用 accelerate 的显存管理
+os.environ["ACCELERATE_USE_MEMORY_EFFICIENT_ATTENTION"] = "false"
+os.environ["ACCELERATE_DISABLE_MEMORY_EFFICIENT_ATTENTION"] = "1"
 
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 
@@ -57,123 +59,74 @@ class FluxInterface:
         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"
-            )
-            
-            # 再次清理显存
-            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
+        
+        # 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").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)
+        
+        # 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='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)
+        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
+        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,
+            'text_encoder_two': text_encoder_two,
+            'tokenizer_two': tokenizer_two,
+            'vae': vae,
+            'transformer': transformer,
+            'scheduler': scheduler,
+            'qwen2vl': qwen2vl,
+            'connector': connector
+        }
+        
+        # Initialize processor and 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,
+        ) 
     
     def resize_image(self, img, max_pixels=1050000):
         if not isinstance(img, Image.Image):