Update app.py

app.py

@@ -5,37 +5,35 @@ import gradio as gr
 from dotenv import load_dotenv
 import torch
 from scipy.io.wavfile import write
-from diffusers import DiffusionPipeline  
+from diffusers import DiffusionPipeline
 from transformers import pipeline
 from pathlib import Path
 
 load_dotenv()
 hf_token = os.getenv("HF_TKN")
 
-# Initialize pipelines globally (in CPU mode)
+device_id = 0 if torch.cuda.is_available() else -1
+
 captioning_pipeline = pipeline(
     "image-to-text",
-    model="nlpconnect/vit-gpt2-image-captioning"
+    model="nlpconnect/vit-gpt2-image-captioning",
+    device=device_id
 )
 
 pipe = DiffusionPipeline.from_pretrained(
     "cvssp/audioldm2",
     use_auth_token=hf_token
 )
+# The AudioLDM pipeline can be moved to CUDA/CPU explicitly inside the function.
 
 @spaces.GPU(duration=120)
 def analyze_image_with_free_model(image_file):
     try:
-        # Move captioning pipeline to GPU
-        captioning_pipeline.to("cuda")
         with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
             temp_file.write(image_file)
             temp_image_path = temp_file.name
 
         results = captioning_pipeline(temp_image_path)
-        # Move back to CPU (optional)
-        captioning_pipeline.to("cpu")
-
         if not results or not isinstance(results, list):
             return "Error: Could not generate caption.", True
         
@@ -50,7 +48,6 @@ def analyze_image_with_free_model(image_file):
 @spaces.GPU(duration=120)
 def get_audioldm_from_caption(caption):
     try:
-        # Move AudioLDM pipeline to GPU
         pipe.to("cuda")
         audio_output = pipe(
             prompt=caption,
@@ -61,7 +58,7 @@ def get_audioldm_from_caption(caption):
         audio = audio_output.audios[0]
 
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
-            write(temp_wav.name, 16000, audio) 
+            write(temp_wav.name, 16000, audio)
             return temp_wav.name
 
     except Exception as e:
@@ -78,12 +75,10 @@ css = """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.HTML("""
-    <h1 style="text-align: center;">
-           🎶 Generate Sound Effects from Image
-        </h1>
-         <p style="text-align: center;">
-          ⚡ Powered by <a href="https://bilsimaging.com" target="_blank">Bilsimaging</a>
-        </p>
+    <h1 style="text-align: center;">🎶 Generate Sound Effects from Image</h1>
+    <p style="text-align: center;">
+        ⚡ Powered by <a href="https://bilsimaging.com" target="_blank">Bilsimaging</a>
+    </p>
         """)
 
     gr.Markdown("""