Update app.py

app.py

@@ -1,39 +1,31 @@
+import torch
+from diffusers import FluxPipeline
 import gradio as gr
-from transformers import DalleBartProcessor, FlaxDalleBartForConditionalGeneration
-from PIL import Image
-import numpy as np
-import jax
-import jax.numpy as jnp
 
-# Load the DALL-E Mega model and processor
-processor = DalleBartProcessor.from_pretrained("dalle-mini/dalle-mega")
-model = FlaxDalleBartForConditionalGeneration.from_pretrained("dalle-mini/dalle-mega")
+# Initialize the model
+pipe = FluxPipeline.from_pretrained("Shakker-Labs/AWPortrait-FL", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
 
-# Function to generate an image from a text prompt
 def generate_image(prompt):
-    # Process the prompt
-    inputs = processor([prompt], return_tensors="jax", padding="max_length", truncation=True, max_length=64)
-    
-    # Generate the images
-    outputs = model.generate(**inputs, do_sample=True, num_beams=4, num_return_sequences=1)
-    
-    # Decode the images and convert them to displayable format
-    images = model.decode(outputs.sequences)
-    images = jax.device_get(images)
-    
-    # Convert to a PIL image
-    pil_img = Image.fromarray(np.asarray(images[0]).astype(np.uint8))
-    
-    return pil_img
+    # Generate the image
+    image = pipe(prompt, 
+                 num_inference_steps=24, 
+                 guidance_scale=3.5,
+                 width=768, height=1024,
+                ).images[0]
+    # Save the image
+    image_path = "generated_image.png"
+    image.save(image_path)
+    return image_path
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=generate_image,               # Function to generate the image
-    inputs=gr.Textbox(lines=2, placeholder="Enter your text prompt"),   # Input textbox for the prompt
-    outputs="image",                 # Output as an image
-    title="DALL-E Mini/Mega Image Generator",
-    description="Generate images from text prompts using the DALL-E Mega model."
+# Define the Gradio interface
+interface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+    outputs=gr.Image(type="file", label="Generated Image"),
+    title="Image Generator",
+    description="Generate images based on the given prompt using the FluxPipeline model."
 )
 
-# Launch the app
-iface.launch()
+# Launch the Gradio app
+interface.launch()