Create app.py

app.py

@@ -0,0 +1,45 @@
+import gradio as gr
+from diffusers import DiffusionPipeline
+from PIL import Image
+import torch
+
+# Load the model and LoRA weights
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Assuming you're using a GPU
+pipe.load_lora_weights("enhanceaiteam/Flux-Uncensored-V2")
+
+# Function for text-to-image
+def text_to_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Function for image-to-image
+def image_to_image(init_image, prompt):
+    init_image = init_image.convert("RGB")  # Ensuring image is RGB
+    image = pipe(prompt, image=init_image).images[0]  # Passing the image as the initial image for editing
+    return image
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# Flux Uncensored V2 Demo")
+
+    with gr.Tab("Text-to-Image"):
+        with gr.Row():
+            text_prompt = gr.Textbox(label="Enter your prompt")
+            generated_image = gr.Image(label="Generated Image")
+        generate_button = gr.Button("Generate")
+        
+        # Connect the button to text-to-image function
+        generate_button.click(fn=text_to_image, inputs=text_prompt, outputs=generated_image)
+
+    with gr.Tab("Image-to-Image"):
+        with gr.Row():
+            init_image = gr.Image(source="upload", label="Upload Initial Image", type="pil")
+            edit_prompt = gr.Textbox(label="Enter your prompt")
+        edited_image = gr.Image(label="Edited Image")
+        edit_button = gr.Button("Generate from Image")
+        
+        # Connect the button to image-to-image function
+        edit_button.click(fn=image_to_image, inputs=[init_image, edit_prompt], outputs=edited_image)
+
+demo.launch()