Update app.py

app.py

@@ -1,10 +1,9 @@
-import torch
 from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
-from diffusers import FluxPipeline
+from diffusers import StableDiffusionPipeline
+import torch
 import gradio as gr
-import io
-from PIL import Image
 import os
+from PIL import Image
 
 # Load the Hugging Face API key from environment variables
 hf_api_key = os.getenv("multimodel_token")
@@ -22,15 +21,15 @@ 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 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
+# Load the alvdansen/flux-koda image generation model using Diffusers
+pipe = StableDiffusionPipeline.from_pretrained("alvdansen/flux-koda", use_auth_token=hf_api_key)
+pipe.to("cuda")  # Use GPU for faster generation, if available
 
-# Function to generate an image using FLUX.1-dev model
+# Function to generate an image using alvdansen/flux-koda 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
+        # Generate the image using the alvdansen/flux-koda model
         image = pipe(translated_text).images[0]
         print("Image generation completed.")
         return image, None