Update app.py

app.py

@@ -4,11 +4,11 @@ from diffusers import DiffusionPipeline
 
 
 # Function to generate image based on input text
-def generate_image(prompt):
+def generate_image(basem, model, prompt):
     # Load the pipeline
     
-    pipeline = DiffusionPipeline.from_pretrained("John6666/mala-anime-mix-nsfw-pony-xl-v3-sdxl")
-    pipeline.load_lora_weights("nevreal/vMurderDrones")
+    pipeline = DiffusionPipeline.from_pretrained(basem)
+    pipeline.load_lora_weights(model)
 
     # Generate the image using the text prompt
     image = pipeline(prompt).images[0]
@@ -18,19 +18,25 @@ def generate_image(prompt):
 with gr.Blocks() as demo:
     # Title
     gr.Markdown("# Text-to-Image Generation WebUI")
+
+    with gr.Row():
+        base_model = gr.Textbox(label="Enter your base model here", placeholder="John6666/mala-anime-mix-nsfw-pony-xl-v3-sdxl")
+        main_model = gr.Textbox(label="Enter your main model here", placeholder="nevreal/vMurderDrones")
+    
     
     # Input for text prompt
     with gr.Row():
         prompt = gr.Textbox(label="Enter your prompt here", placeholder="Type your text prompt...")
     
     # Output image display
-    output_image = gr.Image(label="Generated Image")
+    with gr.Row(scale=4):
+        output_image = gr.Image(label="Generated Image")
 
     # Button to trigger the image generation
     generate_button = gr.Button("Generate Image")
 
     # When the button is clicked, call the generate_image function
-    generate_button.click(fn=generate_image, inputs=prompt, outputs=output_image)
+    generate_button.click(fn=generate_image, inputs=base_model, main_model, prompt, outputs=output_image)
 
 # Launch the interface
 demo.launch()