Update app.py

app.py

@@ -2,90 +2,75 @@ import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
-# Load the base Stable Diffusion model (v1.5) in half-precision for efficiency
+# Load the base model
 model_id = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float16, 
-    safety_checker=None   # disable safety checker for faster inference (floor plans are non-NSFW)
+    model_id,
+    torch_dtype=torch.float16,
+    safety_checker=None
 )
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe.enable_attention_slicing()  # reduce memory usage, good for running on limited GPU
+pipe.enable_attention_slicing()
 
-# Move the pipeline to GPU (if available)
+# Move pipeline to GPU or CPU
 pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load the FloorPlan LoRA weights into the pipeline
+# Load FloorPlan LoRA weights
 pipe.load_lora_weights(
-    "ejazhabibdar/sd-FloorPlan-model",  # Hugging Face model hub ID for the LoRA weights
-    weight_name="pytorch_lora_weights.safetensors"  # file containing LoRA weights
-):contentReference[oaicite:1]{index=1}
-
-# (Optional) Enable xFormers for faster/less memory if installed
-# try:
-#     pipe.enable_xformers_memory_efficient_attention()
-# except Exception as e:
-#     pass
+    "ejazhabibdar/sd-FloorPlan-model",
+    weight_name="pytorch_lora_weights.safetensors"
+)
 
-# Define a function to generate the floor plan image based on user inputs
+# Prompt generation logic
 def generate_floorplan(bedrooms, bathrooms, floors, plot_size):
-    # Construct a text prompt from the structured inputs:
-    # We use the LoRA's recommended style: "Floor plan of a [small/big] house, few/many rooms, one/multiple bathrooms, ...":contentReference[oaicite:2]{index=2}
     prompt_parts = []
-    # Determine house size descriptor
-    if plot_size.lower() in ["small", "medium", "large"]:
-        size_word = "small" if plot_size.lower() == "small" else ("big" if plot_size.lower() == "large" else "")
-        if size_word:
-            prompt_parts.append(f"{size_word} house")
-        else:
-            prompt_parts.append("house")
+
+    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")
-    # Rooms descriptor based on bedrooms count
+
     if bedrooms <= 2:
         prompt_parts.append("few rooms")
     elif bedrooms >= 5:
         prompt_parts.append("many rooms")
     else:
         prompt_parts.append(f"{bedrooms} rooms")
-    # Bathrooms descriptor
+
     if bathrooms == 1:
         prompt_parts.append("one bathroom")
     else:
         prompt_parts.append(f"{bathrooms} bathrooms")
-    # Floors descriptor
+
     if floors == 1:
-        # single floor might not need explicit mention
         prompt_parts.append("one floor")
     else:
         prompt_parts.append(f"{floors} floors")
-    # You can add more prompt details (kitchen size, windows) if desired, but we'll keep it simple
+
     prompt = "Floor plan of a " + ", ".join(prompt_parts) + "."
-    
-    # Generate the image with the pipeline
-    image = pipe(prompt, 
-                 num_inference_steps=25,   # number of diffusion steps
-                 guidance_scale=7.5        # classifier-free guidance scale
-                 # You can adjust the LoRA strength via cross_attention_kwargs if needed, e.g. {"scale": 0.8}
-                ).images[0]
+
+    image = pipe(prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
     return image
 
-# Build the Gradio interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏠 AI Floor Plan Generator\nEnter your desired house parameters and click **Generate** to create a floor plan:")
+    gr.Markdown("## 🏠 AI Floor Plan Generator")
     with gr.Row():
-        bedrooms = gr.Slider(label="Bedrooms", minimum=1, maximum=10, value=3, step=1)
-        bathrooms = gr.Slider(label="Bathrooms", minimum=1, maximum=5, value=2, step=1)
+        bedrooms = gr.Slider(label="Bedrooms", minimum=1, maximum=10, value=3)
+        bathrooms = gr.Slider(label="Bathrooms", minimum=1, maximum=5, value=2)
     with gr.Row():
-        floors = gr.Slider(label="Floors", minimum=1, maximum=3, value=1, step=1)
+        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")
-    
-    # Connect the input components to the generation function
-    generate_btn.click(fn=generate_floorplan, 
-                       inputs=[bedrooms, bathrooms, floors, plot_size], 
-                       outputs=output_image)
 
-# Launch the app (not strictly required on Spaces, but okay to include)
+    generate_btn.click(
+        fn=generate_floorplan,
+        inputs=[bedrooms, bathrooms, floors, plot_size],
+        outputs=output_image
+    )
+
 demo.launch()