Update app.py

app.py

@@ -28,11 +28,11 @@ except Exception as e:
     print(f"Error during Hugging Face login: {str(e)}")
 
 # Set device: CUDA if available, else CPU
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 model_id1 = os.getenv("API_KEY")
 pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
-pipe = pipe.to("cuda")
+pipe = pipe.to(device)
 
 # Updated function for text generation using the new API structure
 def generate_creative_text(prompt):
@@ -81,7 +81,7 @@ def process_audio(audio_path, image_option, creative_text_option):
         try:
             model_id1 = "dreamlike-art/dreamlike-diffusion-1.0"
             pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
-            pipe = pipe.to("cuda")
+            pipe = pipe.to(device)
             image = pipe(translation).images[0]
         except Exception as e:
             return tamil_text, translation, creative_text, f"An error occurred during image generation: {str(e)}"