Update mini.py

mini.py

@@ -133,10 +133,7 @@ pipeline = FluxPipeline(
     text_encoder=text_encoder,
     tokenizer=tokenizer,
 )
-# Move Qwen2VL models to GPU
-logger.info("Moving Qwen2VL models to GPU...")
-qwen2vl.to(DEVICE)
-connector.to(DEVICE)
+
         
 # # Move Transformer and VAE to GPU
 # logger.info("Moving Transformer and VAE to GPU...")
@@ -150,7 +147,11 @@ connector.to(DEVICE)
 def process_image(image):
     """Process image with Qwen2VL model"""
     try:
-
+        # Move Qwen2VL models to GPU
+        logger.info("Moving Qwen2VL models to GPU...")
+        qwen2vl.to(DEVICE)
+        connector.to(DEVICE)
+        
         message = [
             {
                 "role": "user",
@@ -181,8 +182,8 @@ def process_image(image):
             result = (image_hidden_state.cpu(), image_grid_thw)
         
         # Move models back to CPU
-        # qwen2vl.cpu()
-        # connector.cpu()
+        qwen2vl.cpu()
+        connector.cpu()
         torch.cuda.empty_cache()
         
         return result