Update app.py

app.py

@@ -109,6 +109,9 @@ def generate_image_with_flux(
     # Initialize FLUX pipeline here
     dtype = torch.bfloat16
     device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    torch.cuda.empty_cache()  # Clear cache
+    gc.collect()  # Run garbage collection
     flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
     if randomize_seed:
@@ -151,13 +154,10 @@ def merge_audio_files(mp3_names: List[str]) -> str:
 @spaces.GPU()
 def get_output_video(text, seed, randomize_seed, width, height, num_inference_steps):
     print("DEBUG: Starting get_output_video function...")
-    
     # Set the device here, inside the GPU-accelerated function
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
+    device = "cuda" if torch.cuda.is_available() else "cpu" 
     # Move the model to the GPU
     model.to(device)
-
     # Summarize the input text
     print("DEBUG: Summarizing text...")
     inputs = tokenizer(
@@ -174,9 +174,7 @@ def get_output_video(text, seed, randomize_seed, width, height, num_inference_st
     )
     plot = list(summary[0].split('.'))
     print(f"DEBUG: Summary generated: {plot}")
-
     image_system ="Generate a realistic picture about this: "
-
     # Generate images for each sentence in the plot
     generated_images = []
     for i, senten in enumerate(plot[:-1]):
@@ -197,8 +195,8 @@ def get_output_video(text, seed, randomize_seed, width, height, num_inference_st
         print(f"DEBUG: Image generated and saved to {image_path}")
 
         #del min_dalle_model # No need to delete the model here
-        # torch.cuda.empty_cache() # No need to empty cache here
-        # gc.collect() # No need to collect garbage here
+        #torch.cuda.empty_cache() # No need to empty cache here
+        #gc.collect() # No need to collect garbage here
 
     # Create subtitles from the plot
     sentences = plot[:-1]