Update app.py

app.py

@@ -4,10 +4,11 @@ import gradio as gr
 from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
 
     
-device = 'cpu'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 encoder_checkpoint = 'google/vit-base-patch16-224'
-decoder_checkpoint = 'surajp/gpt2-hindi'
-model_checkpoint = 'team-indain-image-caption/hindi-image-captioning'
+decoder_checkpoint = 'sachin/vit2distilgpt2'
+model_checkpoint = 'nlpconnect/vit-gpt2-image-captioning'
 feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
 model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)