Update app.py

app.py

@@ -18,11 +18,13 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 base_model_id = "runwayml/stable-diffusion-v1-5"
 model_id = "LuyangZ/FloorAI"
 
-controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float32)
+controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype="auto")
+# controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float32)
 controlnet.to(device)
 torch.cuda.empty_cache()
 
-pipeline = StableDiffusionControlNetPipeline.from_pretrained(base_model_id , controlnet=controlnet, torch_dtype=torch.float32)
+# pipeline = StableDiffusionControlNetPipeline.from_pretrained(base_model_id , controlnet=controlnet, torch_dtype=torch.float32)
+pipeline = StableDiffusionControlNetPipeline.from_pretrained(base_model_id , controlnet=controlnet, torch_dtype="auto")
 pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
 
 pipeline = pipeline.to(device)
@@ -95,7 +97,7 @@ def floorplan_generation(outline, num_of_rooms):
     validation_image = n_outline
 
     image_lst = []
-    for i in range(2):
+    for i in range(5):
         seed = randrange(500)
         generator = torch.Generator(device=device).manual_seed(seed)
     
@@ -111,15 +113,19 @@ def floorplan_generation(outline, num_of_rooms):
         image = clean_img(image, mask)
         image_lst.append(image)
         
-    return image_lst[0], image_lst[1]
+    return image_lst[0], image_lst[1], image_lst[2], image_lst[3], image_lst[4]
 
 
 gradio_interface = gradio.Interface(
   fn=floorplan_generation,
   inputs=[gradio.Image(label="Floor Plan Outline, Entrance"),
           gradio.Textbox(type="text", label="number of rooms", placeholder="number of rooms")],
-  outputs=[gradio.Image(label="Generated Floor Plan 1"), gradio.Image(label="Generated Floor Plan 2")],
-  title="floorplan generation")
+  outputs=[gradio.Image(label="Generated Floor Plan 1"), 
+           gradio.Image(label="Generated Floor Plan 2"),
+           gradio.Image(label="Generated Floor Plan 3"),
+           gradio.Image(label="Generated Floor Plan 4"),
+           gradio.Image(label="Generated Floor Plan 5")],
+  title="FloorAI")
 
 
 gradio_interface.queue(max_size=10, status_update_rate="auto")