ZeroGPU

app.py

@@ -41,27 +41,29 @@ if not os.path.exists(MODEL_CACHE_DIR):
         logger.error(f"Error downloading models: {str(e)}")
         raise
 
-# 加载所有模型到全局变量
-logger.info("Loading models...")
+# 加载小模型到 GPU
+logger.info("Loading small models to 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(dtype)
+).to(dtype).to(device)
 
 text_encoder_two = T5EncoderModel.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")
-).to(dtype)
+).to(dtype).to(device)
 
 tokenizer_two = T5TokenizerFast.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
 
+# 大模型初始加载到 CPU
+logger.info("Loading large models to CPU...")
 vae = AutoencoderKL.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "flux/vae")
-).to(dtype)
+).to(dtype).cpu()
 
 transformer = FluxTransformer2DModel.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "flux/transformer")
-).to(dtype)
+).to(dtype).cpu()
 
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "flux/scheduler"),
@@ -70,7 +72,7 @@ scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
 
 qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
     os.path.join(MODEL_CACHE_DIR, "qwen2-vl")
-).to(dtype)
+).to(dtype).cpu()
 
 qwen2vl_processor = AutoProcessor.from_pretrained(
     MODEL_ID,
@@ -79,20 +81,20 @@ qwen2vl_processor = AutoProcessor.from_pretrained(
     max_pixels=256*28*28
 )
 
-# 加载connector和embedder
-connector = nn.Linear(3584, 4096).to(dtype)
+# 加载 connector 和 embedder 到 CPU
+connector = nn.Linear(3584, 4096).to(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.replace('module.', ''): v.to(dtype) for k, v in connector_state.items()}
 connector.load_state_dict(connector_state)
 
-t5_context_embedder = nn.Linear(4096, 3072).to(dtype)
+t5_context_embedder = nn.Linear(4096, 3072).to(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(dtype) for k, v in t5_embedder_state.items()}
 t5_context_embedder.load_state_dict(t5_embedder_state)
 
-# 创建pipeline
+# 创建pipeline (先用CPU上的模型)
 pipeline = FluxPipeline(
     transformer=transformer,
     scheduler=scheduler,
@@ -120,6 +122,11 @@ ASPECT_RATIOS = {
 def process_image(image):
     """Process image with Qwen2VL model"""
     try:
+        # 将 Qwen2VL 相关模型移到 GPU
+        logger.info("Moving Qwen2VL models to GPU...")
+        qwen2vl.to(device)
+        connector.to(device)
+
         message = [
             {
                 "role": "user",
@@ -147,7 +154,16 @@ def process_image(image):
             image_hidden_state = output_hidden_state[image_token_mask].view(1, -1, output_hidden_state.size(-1))
             image_hidden_state = connector(image_hidden_state)
             
-            return (image_hidden_state, image_grid_thw)
+            # 保存结果到 CPU
+            result = (image_hidden_state.cpu(), image_grid_thw)
+            
+            # 将模型移回 CPU 并清理显存
+            logger.info("Moving Qwen2VL models back to CPU...")
+            qwen2vl.cpu()
+            connector.cpu()
+            torch.cuda.empty_cache()
+            
+            return result
             
     except Exception as e:
         logger.error(f"Error in process_image: {str(e)}")
@@ -167,8 +183,14 @@ def compute_t5_text_embeddings(prompt):
     ).to(device)
     
     prompt_embeds = text_encoder_two(text_inputs.input_ids)[0]
+    
+    # 将 t5_context_embedder 移到 GPU
+    t5_context_embedder.to(device)
     prompt_embeds = t5_context_embedder(prompt_embeds)
     
+    # 将 t5_context_embedder 移回 CPU
+    t5_context_embedder.cpu()
+    
     return prompt_embeds
 
 def compute_text_embeddings(prompt=""):
@@ -216,8 +238,18 @@ def generate_images(input_image, prompt="", guidance_scale=3.5,
         # Generate images
         try:
             logger.info("Starting image generation...")
+            
+            # 将 Transformer 和 VAE 移到 GPU
+            logger.info("Moving Transformer and VAE to GPU...")
+            transformer.to(device)
+            vae.to(device)
+            
+            # 更新 pipeline 中的模型引用
+            pipeline.transformer = transformer
+            pipeline.vae = vae
+            
             output_images = pipeline(
-                prompt_embeds=qwen2_hidden_state.repeat(num_images, 1, 1),
+                prompt_embeds=qwen2_hidden_state.to(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,
@@ -225,8 +257,15 @@ def generate_images(input_image, prompt="", guidance_scale=3.5,
                 height=height,
                 width=width,
             ).images
+            
             logger.info("Image generation completed")
             
+            # 将 Transformer 和 VAE 移回 CPU
+            logger.info("Moving models back to CPU...")
+            transformer.cpu()
+            vae.cpu()
+            torch.cuda.empty_cache()
+            
             return output_images
             
         except Exception as e: