import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("sd-dreambooth-library/herge-style", torch_dtype=torch.float16).to("cuda")

def generate_image(prompt):
    pipeline.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
    image = pipeline(prompt).images[0]
    return image.permute(1, 2, 0).cpu().numpy()

title = "DreamBooth Gradio App"
description = "Generate images based on prompts using the DreamBooth model."

prompt_textbox = gr.Textbox(lines=3, label="Prompt")
image_output = gr.Image(draw_fn=generate_image, label="Generated Image")

gr.Interface(fn=generate_image, inputs=prompt_textbox, outputs=image_output, title=title, description=description).launch()