Update app.py

app.py

@@ -25,19 +25,24 @@ def load_css():
         unsafe_allow_html=True,
     )
 
-@st.cache_resource(show_spinner=False)
 def load_all_models():
+    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+
     model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-    text_gen = pipeline("text-generation", model="gpt2", device=-1)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, trust_remote_code=True)
+
+    text_gen = pipeline("text-generation", model="gpt2")
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-base",
-        torch_dtype=torch.float32,
-        safety_checker=None
-    ).to("cpu")
+        "CompVis/stable-diffusion-v1-4", 
+        revision="fp16", 
+        torch_dtype=torch.float16
+    ).to("cuda")
+
     return tokenizer, model, text_gen, img_pipe
 
+
 def translate_tamil(text, tokenizer, model):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
     outs = model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)