Update app.py

app.py

@@ -3,42 +3,57 @@ from diffusers import StableDiffusionPipeline
 import torch
 
 # Function to automatically switch between GPU and CPU
-def load_model(model_id):
+def load_model(base_model_id, adapter_model_id=None):
     if torch.cuda.is_available():
         device = "cuda"
         info = "Running on GPU (CUDA)"
     else:
         device = "cpu"
         info = "Running on CPU"
-    
-    # Load the model dynamically on the correct device
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+
+    # Load the base model dynamically on the correct device
+    pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
     pipe = pipe.to(device)
+
+    # If an adapter model is provided, load and merge the adapter model
+    if adapter_model_id:
+        adapter_model = StableDiffusionPipeline.from_pretrained(adapter_model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+        pipe.unet.load_attn_procs(adapter_model_id)  # This applies the adapter like LoRA to the model's UNet
+        info += f" with Adapter Model: {adapter_model_id}"
     
     return pipe, info
 
+
+if torch.cuda.is_available():
+        device = "cuda"
+        info = "Running on GPU (CUDA) 🔥"
+    else:
+        device = "cpu"
+        info = "Running on CPU 🥶"
+
 # Function for text-to-image generation with dynamic model ID and device info
-def generate_image(model_id, prompt):
-    pipe, info = load_model(model_id)
+def generate_image(base_model_id, adapter_model_id, prompt):
+    pipe, info = load_model(base_model_id, adapter_model_id)
     image = pipe(prompt).images[0]
     return image, info
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Custom Text-to-Image Generator")
-
+    gr.Markdown("## Custom Text-to-Image Generator with Adapter Support")
+    gr.Markdown(f"{info}")
     with gr.Row():
         with gr.Column():
-            model_id = gr.Textbox(label="Enter Model ID (e.g., nevreal/vMurderDrones)", placeholder="Model ID")
+            base_model_id = gr.Textbox(label="Enter Base Model ID (e.g., CompVis/stable-diffusion-v1-4)", placeholder="Base Model ID")
+            adapter_model_id = gr.Textbox(label="Enter Adapter Model ID (optional, e.g., nevreal/vMurderDrones-Lora)", placeholder="Adapter Model ID (optional)", value="")
             prompt = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to generate")
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column():
             output_image = gr.Image(label="Generated Image")
-            device_info = gr.Markdown()  # To display if GPU or CPU is used
+            device_info = gr.Markdown()  # To display if GPU or CPU is used and whether an adapter is applied
     
     # Link the button to the image generation function
-    generate_btn.click(fn=generate_image, inputs=[model_id, prompt], outputs=[output_image, device_info])
+    generate_btn.click(fn=generate_image, inputs=[base_model_id, adapter_model_id, prompt], outputs=[output_image, device_info])
 
 # Launch the app
 demo.launch()