Update app.py

app.py

@@ -77,24 +77,24 @@ def generate(
     if torch.cuda.is_available():
 
         if not use_img2img:
-            pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+            pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                         
             if use_vae:  
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                 
         if use_img2img:
-            pipe = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+            pipe = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
 
             init_image = load_image(url)
 
             if use_vae:  
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
 
         if use_controlnet:
-            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
-            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16)
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
 
             image = load_image(controlnet_img)
             
@@ -105,12 +105,12 @@ def generate(
             image = Image.fromarray(image)
 
             if use_vae:
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                 
         if use_controlnetinpaint:
-            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
-            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16)
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
 
             image_start = load_image(controlnet_img)
             image = load_image(controlnet_img)
@@ -123,8 +123,8 @@ def generate(
             image = Image.fromarray(image)
             
             if use_vae:
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                 
         if use_lora:
             pipe.load_lora_weights(lora, adapter_name="1")
@@ -474,4 +474,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20, default_concurrency_limit=5).launch()
+    demo.queue(max_size=20, default_concurrency_limit=2).launch()