revamped code

image_generator.py

@@ -31,34 +31,21 @@
 #     print("Generated Image Path:", generate_image(prompt))
 
 
-# from diffusers import StableDiffusionPipeline
-# import torch
-
-# # Preload the model globally
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-# model = StableDiffusionPipeline.from_pretrained(
-#     "stabilityai/stable-diffusion-2-1-base",
-#     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-# )
-# model.to(device)
-
-# def generate_image(prompt):
-#     """Generate an image from a text prompt."""
-#     image = model(prompt).images[0]
-#     output_path = "output.png"
-#     image.save(output_path)
-#     return output_path
-
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the Stable Diffusion pipeline on the CPU
-pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
-pipe.to("cpu")  # Force the pipeline to run on the CPU
+# Preload the model globally
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5"
+)
+pipeline = pipeline.to(device)
 
 def generate_image(prompt):
-    """Generate an image from a text prompt."""
-    image = pipe(prompt).images[0]
-    image_path = "output_image.png"
+    """Generate an image based on the input prompt."""
+    with torch.no_grad():
+        image = pipeline(prompt).images[0]
+    # Save the image locally and return the file path
+    image_path = "generated_image.png"
     image.save(image_path)
     return image_path

translate.py

@@ -25,13 +25,14 @@
 #     inputs = tokenizer(text, return_tensors="pt", padding=True)
 #     translated = translation_model.generate(**inputs)
 #     return tokenizer.decode(translated[0], skip_special_tokens=True)
-
+import torch
 from transformers import MarianMTModel, MarianTokenizer
 
 # Preload the translation model globally
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
-translation_model = MarianMTModel.from_pretrained(model_name)
+translation_model = MarianMTModel.from_pretrained(model_name).to(device)
 
 # Define supported languages for the dropdown
 SUPPORTED_LANGUAGES = {
@@ -75,7 +76,7 @@ def translate_text(text, src_lang, tgt_lang="en"):
     translation_model = MarianMTModel.from_pretrained(model_name)
 
     # Tokenize the input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True)
+    inputs = tokenizer(text, return_tensors="pt", padding=True).to(device)
     
     # Perform translation
     translated = translation_model.generate(**inputs)