Update app.py

app.py

@@ -2,17 +2,47 @@ import gradio as gr
 from diffusers import AutoPipelineForText2Image
 import torch
 
+# Load Dreambooth model
 pipeline = AutoPipelineForText2Image.from_pretrained("sd-dreambooth-library/herge-style", torch_dtype=torch.float16).to("cuda")
 
 def generate_image(prompt):
+    # Generate image based on 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()
+    return image
 
-title = "DreamBooth Gradio App"
-description = "Generate images based on prompts using the DreamBooth model."
+def image_to_base64(image):
+    # Convert image to base64
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode()
 
-prompt_textbox = gr.Textbox(lines=3, label="Prompt")
-image_output = gr.Image(draw_fn=generate_image, label="Generated Image")
+def base64_to_image(base64_str):
+    # Convert base64 to image
+    image_data = base64.b64decode(base64_str)
+    return Image.open(BytesIO(image_data))
 
-gr.Interface(fn=generate_image, inputs=prompt_textbox, outputs=image_output, title=title, description=description).launch()
+def handle_prompt_image(prompt):
+    # Generate image based on prompt and convert to base64
+    image = generate_image(prompt)
+    base64_str = image_to_base64(image)
+    return base64_str
+
+def main():
+    # Interface setup
+    image_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here...")
+    prompt_output = gr.Textbox(label="Base64 Encoded Image", readonly=True)
+
+    iface = gr.Interface(
+        fn=handle_prompt_image,
+        inputs=image_input,
+        outputs=prompt_output,
+        title="Dreambooth Image Generator",
+        description="Enter a prompt to generate an image using the Dreambooth model.",
+        theme="compact"
+    )
+
+    iface.launch(share=True)
+
+if __name__ == "__main__":
+    main()