Optimize memory

app.py

@@ -643,57 +643,63 @@ def worker_start_end(input_image, end_image, image_position, prompts, n_prompt,
 
         H, W, C = input_image.shape
         height, width = find_nearest_bucket(H, W, resolution=resolution)
-        input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
-
-        Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}_start.png'))
-
-        input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
-        input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
         
-        # Processing end image (if provided)
-        has_end_image = end_image is not None
-        if has_end_image:
-            stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Processing end frame ...'))))
+        def get_start_latent(input_image, end_image, height, width, vae, gpu, image_encoder, high_vram):
+            input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
             
-            H_end, W_end, C_end = end_image.shape
-            end_image_np = resize_and_center_crop(end_image, target_width=width, target_height=height)
+            Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}_start.png'))
             
-            Image.fromarray(end_image_np).save(os.path.join(outputs_folder, f'{job_id}_end.png'))
+            input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
+            input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
             
-            end_image_pt = torch.from_numpy(end_image_np).float() / 127.5 - 1
-            end_image_pt = end_image_pt.permute(2, 0, 1)[None, :, None]
-
-        # VAE encoding
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
-
-        if not high_vram:
-            load_model_as_complete(vae, target_device=gpu)
-
-        start_latent = vae_encode(input_image_pt, vae)
-        
-        if has_end_image:
-            end_latent = vae_encode(end_image_pt, vae)
-
-        # CLIP Vision
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
-
-        if not high_vram:
-            load_model_as_complete(image_encoder, target_device=gpu)
-
-        image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
-        image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
-        
-        if has_end_image:
-            end_image_encoder_output = hf_clip_vision_encode(end_image_np, feature_extractor, image_encoder)
-            end_image_encoder_last_hidden_state = end_image_encoder_output.last_hidden_state
-            # Combine both image embeddings or use a weighted approach
-            image_encoder_last_hidden_state = (image_encoder_last_hidden_state + end_image_encoder_last_hidden_state) / 2
-
-        # Clean GPU
-        if not high_vram:
-            unload_complete_models(
-                image_encoder
-            )
+            # Processing end image (if provided)
+            has_end_image = end_image is not None
+            if has_end_image:
+                stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Processing end frame ...'))))
+                
+                H_end, W_end, C_end = end_image.shape
+                end_image_np = resize_and_center_crop(end_image, target_width=width, target_height=height)
+                
+                Image.fromarray(end_image_np).save(os.path.join(outputs_folder, f'{job_id}_end.png'))
+                
+                end_image_pt = torch.from_numpy(end_image_np).float() / 127.5 - 1
+                end_image_pt = end_image_pt.permute(2, 0, 1)[None, :, None]
+            
+            # VAE encoding
+            stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
+            
+            if not high_vram:
+                load_model_as_complete(vae, target_device=gpu)
+            
+            start_latent = vae_encode(input_image_pt, vae)
+            
+            if has_end_image:
+                end_latent = vae_encode(end_image_pt, vae)
+            
+            # CLIP Vision
+            stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
+            
+            if not high_vram:
+                load_model_as_complete(image_encoder, target_device=gpu)
+            
+            image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
+            image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
+            
+            if has_end_image:
+                end_image_encoder_output = hf_clip_vision_encode(end_image_np, feature_extractor, image_encoder)
+                end_image_encoder_last_hidden_state = end_image_encoder_output.last_hidden_state
+                # Combine both image embeddings or use a weighted approach
+                image_encoder_last_hidden_state = (image_encoder_last_hidden_state + end_image_encoder_last_hidden_state) / 2
+            
+            # Clean GPU
+            if not high_vram:
+                unload_complete_models(
+                    image_encoder
+                )
+                
+            return [start_latent, image_encoder_last_hidden_state]
+            
+        [start_latent, image_encoder_last_hidden_state] = get_start_latent(input_image, end_image, height, width, vae, gpu, image_encoder, high_vram)
 
         # Dtype
         image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)