app.py

import torch
from diffusers import FluxPipeline
import gradio as gr

# Initialize the model
pipe = FluxPipeline.from_pretrained("Shakker-Labs/AWPortrait-FL", torch_dtype=torch.bfloat16)
pipe.to("cuda")

def generate_image(prompt):
    # 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

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