Update app.py

app.py

@@ -3,28 +3,40 @@ import wget
 from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 import torch
+import time
 
 # Define the device to use (either "cuda" for GPU or "cpu" for CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the models
-caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
-sd_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
+# Function to load the models
+def load_models():
+    global caption_image, sd_pipeline, translator
+    start_time = time.time()
 
-# Load the translation model (English to Arabic)
-translator = pipeline(
-    task="translation",
-    model="facebook/nllb-200-distilled-600M",
-    torch_dtype=torch.bfloat16,
-    device=device
-)
+    # Load the image captioning model
+    caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
+    print(f"Caption model loaded in {time.time() - start_time:.2f} seconds")
+
+    # Load the Stable Diffusion model
+    sd_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
+    print(f"Stable Diffusion model loaded in {time.time() - start_time:.2f} seconds")
+
+    # Load the translation model (English to Arabic)
+    translator = pipeline(
+        task="translation",
+        model="facebook/nllb-200-distilled-600M",
+        torch_dtype=torch.bfloat16,
+        device=device
+    )
+    print(f"Translator model loaded in {time.time() - start_time:.2f} seconds")
+
+# Load the models
+load_models()
 
 # Download the images
 url1 = "https://github.com/Shahad-b/Image-database/blob/main/sea.jpg?raw=true"
 sea = wget.download(url1)
 
-
-
 # Function to generate images based on the image's caption
 def generate_image_and_translate(image, num_images=1):
     # Generate caption in English from the uploaded image