Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from gradio_client import Client, handle_file
 import torch
 import spaces
-from diffusers import Lumina2Pipeline
+from diffusers import FluxPipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 if torch.cuda.is_available():
@@ -20,8 +20,8 @@ def set_client_for_session(request: gr.Request):
 
 # Load models
 def load_models():    
-    pipe = Lumina2Pipeline.from_pretrained(
-        "X-ART/LeX-Lumina", 
+    pipe = FluxPipeline.from_pretrained(
+        "X-ART/LeX-FLUX", 
         torch_dtype=torch.bfloat16
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -57,18 +57,16 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
     print(f"enhanced caption:\n{enhanced_caption}")
     
     generator = torch.Generator("cpu").manual_seed(seed) if seed != 0 else None
-    
+
     image = pipe(
         enhanced_caption,
         height=1024,
         width=1024,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        cfg_trunc_ratio=1,
-        cfg_normalization=True,
-        max_sequence_length=256,
-        generator=generator,
-        system_prompt="You are an assistant designed to generate superior images with the superior degree of image-text alignment based on textual prompts or user prompts.",
+        guidance_scale=3.5,
+        output_type="pil",
+        num_inference_steps=28,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
     ).images[0]
 
     print(image)