Initial Commit

app.py

@@ -91,34 +91,36 @@ class FluxInterface:
         tokenizer_two = T5TokenizerFast.from_pretrained(
             os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
         
-        # 2. 将大模型初始加载到CPU
+        # 2. 将大模型加载到CPU，但保持bfloat16精度
         vae = AutoencoderKL.from_pretrained(
             os.path.join(MODEL_CACHE_DIR, "flux/vae")
-        ).to(torch.float32).cpu()
+        ).to(self.dtype).cpu()
         
         transformer = FluxTransformer2DModel.from_pretrained(
             os.path.join(MODEL_CACHE_DIR, "flux/transformer")
-        ).to(torch.float32).cpu()
+        ).to(self.dtype).cpu()
         
         scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/scheduler"), 
+            os.path.join(MODEL_CACHE_DIR, "flux/scheduler"),
             shift=1
         )
         
-        # 3. Qwen2VL初始加载到CPU
+        # 3. Qwen2VL加载到CPU，保持bfloat16
         qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
             os.path.join(MODEL_CACHE_DIR, "qwen2-vl")
-        ).to(torch.float32).cpu()
+        ).to(self.dtype).cpu()
         
-        # 4. 加载connector和embedder到CPU
-        connector = Qwen2Connector().to(torch.float32).cpu()
+        # 4. 加载connector和embedder，保持bfloat16
+        connector = Qwen2Connector().to(self.dtype).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)
         
-        self.t5_context_embedder = nn.Linear(4096, 3072).to(torch.float32).cpu()
+        self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype).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)
         
         # 5. 设置所有模型为eval模式
@@ -141,7 +143,6 @@ class FluxInterface:
             'connector': connector
         }
         
-        # Initialize processor and pipeline
         self.qwen2vl_processor = AutoProcessor.from_pretrained(
             self.MODEL_ID,
             subfolder="qwen2-vl",
@@ -160,15 +161,17 @@ class FluxInterface:
     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.to(self.dtype).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)
+            logger.info("Moving Qwen2VL models to GPU...")
+            self.models['qwen2vl'] = self.models['qwen2vl'].to(self.device)
+            self.models['connector'] = self.models['connector'].to(self.device)
+            logger.info("Qwen2VL models moved to GPU")
             
             message = [
                 {
@@ -200,10 +203,12 @@ class FluxInterface:
                 # 保存结果到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')
+            # 2. 将Qwen2VL相关模型移回CPU
+            logger.info("Moving Qwen2VL models back to CPU...")
+            self.models['qwen2vl'] = self.models['qwen2vl'].cpu()
+            self.models['connector'] = self.models['connector'].cpu()
             torch.cuda.empty_cache()
+            logger.info("Qwen2VL models moved to CPU and GPU cache cleared")
             
             return result
             
@@ -242,8 +247,8 @@ class FluxInterface:
         ).to(self.device)
         
         prompt_embeds = self.models['text_encoder_two'](text_inputs.input_ids)[0]
-        prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=self.device)
-        prompt_embeds = self.t5_context_embedder(prompt_embeds)
+        prompt_embeds = self.t5_context_embedder.to(self.device)(prompt_embeds)
+        self.t5_context_embedder = self.t5_context_embedder.cpu()
         
         return prompt_embeds
 
@@ -261,9 +266,9 @@ class FluxInterface:
                 text_inputs.input_ids,
                 output_hidden_states=False
             )
-            pooled_prompt_embeds = prompt_embeds.pooler_output.to(self.dtype)
+            pooled_prompt_embeds = prompt_embeds.pooler_output
+            return pooled_prompt_embeds
 
-        return pooled_prompt_embeds
 
     def generate(self, input_image, prompt="", guidance_scale=3.5, 
                 num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
@@ -275,27 +280,36 @@ class FluxInterface:
                 
             if seed is not None:
                 torch.manual_seed(seed)
+                logger.info(f"Set random seed to: {seed}")
                  
             # 1. 使用Qwen2VL处理图像
+            logger.info("Processing input image with Qwen2VL...")
             qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
+            logger.info("Image processing completed")
             
             # 2. 计算文本嵌入
+            logger.info("Computing text embeddings...")
             pooled_prompt_embeds = self.compute_text_embeddings("")
             t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
+            logger.info("Text embeddings computed")
             
             # 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)
+            logger.info("Moving Transformer and VAE to GPU...")
+            self.models['transformer'] = self.models['transformer'].to(self.device)
+            self.models['vae'] = self.models['vae'].to(self.device)
             
-            # 更新pipeline中的模型
+            # 更新pipeline中的模型引用
             self.pipeline.transformer = self.models['transformer']
             self.pipeline.vae = self.models['vae']
+            logger.info("Models moved to GPU")
             
             # 获取维度
             width, height = ASPECT_RATIOS[aspect_ratio]
+            logger.info(f"Using dimensions: {width}x{height}")
             
             # 4. 生成图像
             try:
+                logger.info("Starting image generation...")
                 output_images = self.pipeline(
                     prompt_embeds=qwen2_hidden_state.to(self.device).repeat(num_images, 1, 1),
                     pooled_prompt_embeds=pooled_prompt_embeds,
@@ -305,11 +319,14 @@ class FluxInterface:
                     height=height,
                     width=width,
                 ).images
+                logger.info("Image generation completed")
                 
                 # 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')
+                logger.info("Moving models back to CPU...")
+                self.models['transformer'] = self.models['transformer'].cpu()
+                self.models['vae'] = self.models['vae'].cpu()
                 torch.cuda.empty_cache()
+                logger.info("Models moved to CPU and GPU cache cleared")
                 
                 return output_images
                 
@@ -323,8 +340,6 @@ class FluxInterface:
 # Initialize the interface
 interface = FluxInterface()
 
-# 直接将 GPU 装饰器应用在最外层的处理函数上
-@spaces.GPU(duration=300)
 def process_request(input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
     """主处理函数，直接处理用户请求"""
     try: