Jesivn commited on
Commit
3266ad7
·
verified ·
1 Parent(s): 6cd7466

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -4,10 +4,11 @@ from diffusers import DiffusionPipeline
4
  import torch
5
  import time
6
 
 
7
  # Load models
8
  translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
9
  summarizer = pipeline("summarization", model="Falconsai/text_summarization")
10
- image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
11
 
12
  # Functions for each task
13
  def translate_tamil_to_english(text):
 
4
  import torch
5
  import time
6
 
7
+ device = "cuda" if torch.cuda.is_available() else "cpu"
8
  # Load models
9
  translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
10
  summarizer = pipeline("summarization", model="Falconsai/text_summarization")
11
+ image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)..to(device)
12
 
13
  # Functions for each task
14
  def translate_tamil_to_english(text):