Update app.py

app.py

@@ -1,24 +1,33 @@
-from diffusers import StableDiffusionPipeline
 import torch
+from diffusers import FluxPipeline
 import gradio as gr
 
-# Load the model (assuming it uses the Stable Diffusion architecture)
-model_id = "Shakker-Labs/AWPortrait-FL"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-pipe.to("cpu")  # Use "cuda" if you have a GPU
+# Load the model and set to CPU
+pipe = FluxPipeline.from_pretrained("Shakker-Labs/AWPortrait-FL", torch_dtype=torch.float32)
+pipe.to("cpu")  # Set to CPU
 
-# Generate an image from a text prompt
+# Define the function to generate an image from a given prompt
 def generate_image(prompt):
-    image = pipe(prompt).images[0]
+    image = pipe(prompt, 
+                 num_inference_steps=24, 
+                 guidance_scale=3.5,
+                 width=768, height=1024).images[0]
+    image.save(f"example.png")
     return image
-    
-# Define the function to generate the image
-def chatbot_image(prompt):
+
+# Create a Gradio interface
+def inference(prompt):
     image = generate_image(prompt)
     return image
 
-# Create a Gradio interface
-interface = gr.Interface(fn=chatbot_image, inputs="text", outputs="image", title="Text-to-Image Chatbot")
+# Gradio Interface
+interface = gr.Interface(
+    fn=inference, 
+    inputs=gr.Textbox(lines=2, placeholder="Enter your image description here..."), 
+    outputs="image",
+    title="Text-to-Image Generator",
+    description="Enter a text prompt to generate an image using AWPortrait-FL."
+)
 
-# Launch the interface
-interface.launch()
+# Launch the Gradio interface
+interface.launch()