Update model.py

model.py

@@ -90,6 +90,7 @@ class Model:
         gc.collect()
         model_id = CONTROLNET_MODEL_IDS[task_name]
         controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
+        print(f'Loaded {model_id}...')
         controlnet.to(self.device)
         torch.cuda.empty_cache()
         gc.collect()
@@ -123,7 +124,7 @@ class Model:
             num_inference_steps=num_steps,
             generator=generator,
             image=control_image,
-        ).images
+        ).images[0]
 
     
     def resize_image(self, image):