Update app.py

app.py

@@ -59,18 +59,30 @@ def generate(
     if torch.cuda.is_available():
             
         if not use_img2img:
-            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16) 
+            pipe.enable_model_cpu_offload()
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            pipe.unet.set_default_attn_processor()
 
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                 pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+                pipe.enable_model_cpu_offload()
+                pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+                pipe.unet.set_default_attn_processor()
 
         if use_img2img:
             pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
-
+            pipe.enable_model_cpu_offload()
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            pipe.unet.set_default_attn_processor()
+            
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                 pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+                pipe.enable_model_cpu_offload()
+                pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+                pipe.unet.set_default_attn_processor()
 
             response = requests.get(url)
             init_image = Image.open(BytesIO(response.content)).convert("RGB")
@@ -82,7 +94,6 @@ def generate(
             
         else:
             pipe.to(device)
-            pipe.unet.set_default_attn_processor()
         
     generator = torch.Generator().manual_seed(seed)