Initial Commit

app.py

@@ -13,7 +13,6 @@ import sys
 
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 from huggingface_hub import snapshot_download
-import spaces
 
 # 设置日志
 logging.basicConfig(
@@ -78,42 +77,53 @@ class FluxInterface:
             return
 
         logger.info("Starting model loading...")
-        # 3. 显式设置 PyTorch 缓存分配器的行为
-        torch.cuda.set_per_process_memory_fraction(0.95)  # 允许使用95%的显存
-        torch.cuda.max_memory_allocated = lambda *args, **kwargs: 0  # 忽略已分配内存的限制
         
-        # Load FLUX components
+        # 1. 首先加载较小的模型到GPU
         tokenizer = CLIPTokenizer.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer"))
-        text_encoder = CLIPTextModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/text_encoder")).to(self.dtype).to(self.device)
-        text_encoder_two = T5EncoderModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")).to(self.dtype).to(self.device)
-        tokenizer_two = T5TokenizerFast.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
+        text_encoder = CLIPTextModel.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/text_encoder")
+        ).to(self.dtype).to(self.device)
         
-        # Load VAE and transformer
-        vae = AutoencoderKL.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/vae")).to(self.dtype).to(self.device)
-        transformer = FluxTransformer2DModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/transformer")).to(self.dtype).to(self.device)
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/scheduler"), shift=1)
+        text_encoder_two = T5EncoderModel.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")
+        ).to(self.dtype).to(self.device)
         
-        # Load Qwen2VL components
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "qwen2-vl")).to(self.dtype).to(self.device)
+        tokenizer_two = T5TokenizerFast.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
         
-        # 加载 connector
-        connector = Qwen2Connector().to(self.dtype).to(self.device)
+        # 2. 将大模型初始加载到CPU
+        vae = AutoencoderKL.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/vae")
+        ).to(torch.float32).cpu()
+        
+        transformer = FluxTransformer2DModel.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/transformer")
+        ).to(torch.float32).cpu()
+        
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "flux/scheduler"), 
+            shift=1
+        )
+        
+        # 3. Qwen2VL初始加载到CPU
+        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
+            os.path.join(MODEL_CACHE_DIR, "qwen2-vl")
+        ).to(torch.float32).cpu()
+        
+        # 4. 加载connector和embedder到CPU
+        connector = Qwen2Connector().to(torch.float32).cpu()
         connector_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/connector.pt")
         connector_state = torch.load(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)
-    
-        # 加载 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(torch.float32).cpu()
         t5_embedder_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/t5_embedder.pt")
         t5_embedder_state = torch.load(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)
         
-        # Set models to eval mode
-        for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder]:
+        # 5. 设置所有模型为eval模式
+        for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, 
+                     connector, self.t5_context_embedder]:
             model.requires_grad_(False)
             model.eval()
         
@@ -133,9 +143,9 @@ class FluxInterface:
         
         # Initialize processor and 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
         )
         
@@ -145,7 +155,61 @@ class FluxInterface:
             vae=vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-        ) 
+        )
+
+    def move_to_device(self, model, device):
+        """Helper function to move model to specified device"""
+        if hasattr(model, 'to'):
+            return model.to(device)
+        return model
+
+    def process_image(self, image):
+        """Process image with Qwen2VL model"""
+        try:
+            # 1. 将Qwen2VL相关模型移到GPU
+            self.models['qwen2vl'] = self.move_to_device(self.models['qwen2vl'], self.device)
+            self.models['connector'] = self.move_to_device(self.models['connector'], self.device)
+            
+            message = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": "Describe this image."},
+                    ]
+                }
+            ]
+            text = self.qwen2vl_processor.apply_chat_template(
+                message, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
+
+            with torch.no_grad():
+                inputs = self.qwen2vl_processor(
+                    text=[text], 
+                    images=[image], 
+                    padding=True, 
+                    return_tensors="pt"
+                ).to(self.device)
+                
+                output_hidden_state, image_token_mask, image_grid_thw = self.models['qwen2vl'](**inputs)
+                image_hidden_state = output_hidden_state[image_token_mask].view(1, -1, output_hidden_state.size(-1))
+                image_hidden_state = self.models['connector'](image_hidden_state)
+                
+                # 保存结果到CPU
+                result = (image_hidden_state.cpu(), image_grid_thw)
+            
+            # 2. 将Qwen2VL相关模型移回CPU以释放显存
+            self.models['qwen2vl'] = self.move_to_device(self.models['qwen2vl'], 'cpu')
+            self.models['connector'] = self.move_to_device(self.models['connector'], 'cpu')
+            torch.cuda.empty_cache()
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"Error in process_image: {str(e)}")
+            raise
     
     def resize_image(self, img, max_pixels=1050000):
         if not isinstance(img, Image.Image):
@@ -163,28 +227,7 @@ class FluxInterface:
             img = img.resize((new_width, new_height), Image.LANCZOS)
         
         return img
-
-    # [Previous methods remain unchanged...]
-    def process_image(self, image):
-        message = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": "Describe this image."},
-                ]
-            }
-        ]
-        text = self.qwen2vl_processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
-
-        with torch.no_grad():
-            inputs = self.qwen2vl_processor(text=[text], images=[image], padding=True, return_tensors="pt").to(self.device)
-            output_hidden_state, image_token_mask, image_grid_thw = self.models['qwen2vl'](**inputs)
-            image_hidden_state = output_hidden_state[image_token_mask].view(1, -1, output_hidden_state.size(-1))
-            image_hidden_state = self.models['connector'](image_hidden_state)
-
-        return image_hidden_state, image_grid_thw
-    
+ 
     def compute_t5_text_embeddings(self, prompt):
         """Compute T5 embeddings for text prompt"""
         if prompt == "":
@@ -222,50 +265,39 @@ class FluxInterface:
 
         return pooled_prompt_embeds
 
-    @spaces.GPU(duration=120)  # 300秒的 GPU 使用时间
-    def generate(self, input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
+    def generate(self, input_image, prompt="", guidance_scale=3.5, 
+                num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
         try:
-            logger.info(f"Starting generation with prompt: {prompt}, guidance_scale: {guidance_scale}, steps: {num_inference_steps}")
+            logger.info(f"Starting generation with prompt: {prompt}")
             
             if input_image is None:
                 raise ValueError("No input image provided")
                 
             if seed is not None:
                 torch.manual_seed(seed)
-                logger.info(f"Set random seed to: {seed}")
-                
-            self.load_models()
-            logger.info("Models loaded successfully")
+                 
+            # 1. 使用Qwen2VL处理图像
+            qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
             
-            # Get dimensions from aspect ratio
-            if aspect_ratio not in ASPECT_RATIOS:
-                raise ValueError(f"Invalid aspect ratio. Choose from {list(ASPECT_RATIOS.keys())}")
-            width, height = ASPECT_RATIOS[aspect_ratio]
-            logger.info(f"Using dimensions: {width}x{height}")
+            # 2. 计算文本嵌入
+            pooled_prompt_embeds = self.compute_text_embeddings("")
+            t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
             
-            # Process input image
-            try:
-                input_image = self.resize_image(input_image)
-                logger.info(f"Input image resized to: {input_image.size}")
-                qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
-                logger.info("Input image processed successfully")
-            except Exception as e:
-                raise RuntimeError(f"Error processing input image: {str(e)}")
+            # 3. 将Transformer和VAE移到GPU
+            self.models['transformer'] = self.move_to_device(self.models['transformer'], self.device)
+            self.models['vae'] = self.move_to_device(self.models['vae'], self.device)
             
-            try:
-                pooled_prompt_embeds = self.compute_text_embeddings("")
-                logger.info("Base text embeddings computed")
-                
-                # Get T5 embeddings if prompt is provided
-                t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
-                logger.info("T5 prompt embeddings computed")
-            except Exception as e:
-                raise RuntimeError(f"Error computing embeddings: {str(e)}")
+            # 更新pipeline中的模型
+            self.pipeline.transformer = self.models['transformer']
+            self.pipeline.vae = self.models['vae']
             
-            # Generate images
+            # 获取维度
+            width, height = ASPECT_RATIOS[aspect_ratio]
+            
+            # 4. 生成图像
             try:
                 output_images = self.pipeline(
-                    prompt_embeds=qwen2_hidden_state.repeat(num_images, 1, 1),
+                    prompt_embeds=qwen2_hidden_state.to(self.device).repeat(num_images, 1, 1),
                     pooled_prompt_embeds=pooled_prompt_embeds,
                     t5_prompt_embeds=t5_prompt_embeds.repeat(num_images, 1, 1) if t5_prompt_embeds is not None else None,
                     num_inference_steps=num_inference_steps,
@@ -274,10 +306,16 @@ class FluxInterface:
                     width=width,
                 ).images
                 
-                logger.info("Images generated successfully")
+                # 5. 将Transformer和VAE移回CPU
+                self.models['transformer'] = self.move_to_device(self.models['transformer'], 'cpu')
+                self.models['vae'] = self.move_to_device(self.models['vae'], 'cpu')
+                torch.cuda.empty_cache()
+                
                 return output_images
+                
             except Exception as e:
                 raise RuntimeError(f"Error generating images: {str(e)}")
+                
         except Exception as e:
             logger.error(f"Error during generation: {str(e)}")
             raise gr.Error(f"Generation failed: {str(e)}")