Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+from diffusers import AutoencoderKL
 from diffusers.utils import load_image
 from controlnet_flux import FluxControlNetModel
 from transformer_flux import FluxTransformer2DModel
@@ -14,10 +15,12 @@ controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Co
 transformer = FluxTransformer2DModel.from_pretrained(
     "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dtype=torch.bfloat16
 )
+vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae").to("cuda")
 pipe = FluxControlNetInpaintingPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     controlnet=controlnet,
     transformer=transformer,
+    vae=vae,
     torch_dtype=torch.bfloat16
 )
 repo_name = "ByteDance/Hyper-SD"
@@ -26,7 +29,7 @@ pipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
 pipe.fuse_lora(lora_scale=0.125)
 pipe.transformer.to(torch.bfloat16)
 pipe.controlnet.to(torch.bfloat16)
-
+pipe.to("cuda")
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     if alignment in ("Left", "Right") and source_width >= target_width:
         return False
@@ -133,7 +136,6 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
 
 @spaces.GPU
 def inpaint(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom, progress=gr.Progress(track_tqdm=True)):
-    pipe.enable_model_cpu_offload()
     
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
     
@@ -158,9 +160,12 @@ def inpaint(image, width, height, overlap_percentage, num_inference_steps, resiz
         controlnet_conditioning_scale=0.9,
         guidance_scale=3.5,
         negative_prompt="",
-        true_guidance_scale=3.5
+        true_guidance_scale=3.5,
+        output_type="latent"
     ).images[0]
-
+    pipe.to("cpu")
+    vae.to("cuda")
+    result = vae.decode(latent_image).sample
     result = result.convert("RGBA")
     cnet_image.paste(result, (0, 0), mask)