app.py

import torch
from diffusers import FluxPipeline
import gradio as gr

# 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

# Define the function to generate an image from a given prompt
def generate_image(prompt):
    image = pipe(prompt, 
                 num_inference_steps=24, 
                 guidance_scale=3.5,
                 width=768, height=1024).images[0]
    image.save(f"example.png")
    return image

# Create a Gradio interface
def inference(prompt):
    image = generate_image(prompt)
    return image

# 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 Gradio interface
interface.launch()