Update app.py

app.py

@@ -1,57 +1,23 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline
 
-# Load the base model
-model_id = "runwayml/stable-diffusion-v1-5"
+# Load base Stable Diffusion model
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
+    "runwayml/stable-diffusion-v1-5",
     torch_dtype=torch.float16,
+    revision="fp16",
     safety_checker=None
 )
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe.enable_attention_slicing()
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Move pipeline to GPU or CPU
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+# Load the FloorPlan LoRA weights
+pipe.load_lora_weights("ejazhabibdar/sd-FloorPlan-model", weight_name="pytorch_lora_weights.safetensors")
 
-# Load FloorPlan LoRA weights
-pipe.load_lora_weights(
-    "ejazhabibdar/sd-FloorPlan-model",
-    weight_name="pytorch_lora_weights.safetensors"
-)
-
-# Prompt generation logic
+# Prompt builder
 def generate_floorplan(bedrooms, bathrooms, floors, plot_size):
-    prompt_parts = []
-
-    size_map = {"small": "small", "medium": "", "large": "big"}
-    size_word = size_map.get(plot_size.lower(), "")
-    if size_word:
-        prompt_parts.append(f"{size_word} house")
-    else:
-        prompt_parts.append("house")
-
-    if bedrooms <= 2:
-        prompt_parts.append("few rooms")
-    elif bedrooms >= 5:
-        prompt_parts.append("many rooms")
-    else:
-        prompt_parts.append(f"{bedrooms} rooms")
-
-    if bathrooms == 1:
-        prompt_parts.append("one bathroom")
-    else:
-        prompt_parts.append(f"{bathrooms} bathrooms")
-
-    if floors == 1:
-        prompt_parts.append("one floor")
-    else:
-        prompt_parts.append(f"{floors} floors")
-
-    prompt = "Floor plan of a " + ", ".join(prompt_parts) + "."
-
-    image = pipe(prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
+    prompt = f"Floor plan of a {plot_size.lower()} house with {bedrooms} bedrooms, {bathrooms} bathrooms, {floors} floors."
+    image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
     return image
 
 # Gradio UI
@@ -63,9 +29,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         floors = gr.Slider(label="Floors", minimum=1, maximum=3, value=1)
         plot_size = gr.Dropdown(label="Plot Size", choices=["Small", "Medium", "Large"], value="Medium")
-
-    generate_btn = gr.Button("Generate Floor Plan")
     output_image = gr.Image(label="Generated Floor Plan", type="pil")
+    generate_btn = gr.Button("Generate Floor Plan")
 
     generate_btn.click(
         fn=generate_floorplan,