Update app.py

app.py

@@ -1,6 +1,6 @@
-from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
-from diffusers import StableDiffusionPipeline
 import torch
+from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
+from diffusers import FluxPipeline
 import gradio as gr
 import io
 from PIL import Image
@@ -22,16 +22,16 @@ text_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
 text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
 text_generator = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)
 
-# Load the FLUX.1-dev image generation model from diffusers
-flux_model = StableDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=hf_api_key)
-flux_model.to("cuda")  # Use GPU for faster generation, if available
+# Load the FLUX.1-dev model using FluxPipeline
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=hf_api_key, torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()  # Save VRAM by offloading the model to CPU
 
 # Function to generate an image using FLUX.1-dev model
 def generate_image_from_text(translated_text):
     try:
         print(f"Generating image from translated text: {translated_text}")
         # Generate the image using the FLUX.1-dev model
-        image = flux_model(translated_text).images[0]
+        image = pipe(translated_text).images[0]
         print("Image generation completed.")
         return image, None
     except Exception as e: