Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 from transformers import pipeline
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, FluxPipeline
 import torch
 import time
 
@@ -8,7 +8,12 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load models
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
 summarizer = pipeline("summarization", model="Falconsai/text_summarization")
-image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
+# image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
+# image_pipe = fluxPipeline("image-generation", model=DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device))
+
+image_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()
+
 
 # Functions for each task
 def translate_tamil_to_english(text):