Update app.py

app.py

@@ -1,35 +1,41 @@
 import gradio as gr
 import torch
-from diffusers import AutoencoderKL
+from diffusers import AutoencoderKL, FluxTransformer2DModel
 from diffusers.utils import load_image
 from controlnet_flux import FluxControlNetModel
 from transformer_flux import FluxTransformer2DModel
 from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
+from transformers import T5EncoderModel, CLIPTextModel
 from PIL import Image, ImageDraw
 import numpy as np
 import spaces
 from huggingface_hub import hf_hub_download
 
+# Load fp8
+#transformer = FluxTransformer2DModel.from_single_file("https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-dev-fp8.safetensors", torch_dtype=torch.bfloat16)
+#quantize(transformer, weights=qfloat8)
+#freeze(transformer)
+
 # Load models
-controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha", torch_dtype=torch.bfloat16)
-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")
-vae.to(torch.bfloat16)
+#controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha", torch_dtype=torch.bfloat16)
+#quantize(controlnet, weights=qfloat8)
+#freeze(controlnet)
+
+text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
+quantize(text_encoder_2, weights=qfloat8)
+freeze(text_encoder_2)
+
 pipe = FluxControlNetInpaintingPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
-    controlnet=controlnet,
-    transformer=transformer,
-    vae=vae,
+    text_encoder_2=None,
     torch_dtype=torch.bfloat16
 )
+pipe.text_encoder_2 = text_encoder_2
+
 repo_name = "ByteDance/Hyper-SD"
 ckpt_name = "Hyper-FLUX.1-dev-8steps-lora.safetensors"
 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:
@@ -150,7 +156,7 @@ def inpaint(image, width, height, overlap_percentage, num_inference_steps, resiz
 
     generator = torch.Generator(device="cuda").manual_seed(42)
 
-    latent_image = pipe(
+    result = pipe(
         prompt=final_prompt,
         height=height,
         width=width,
@@ -162,12 +168,8 @@ def inpaint(image, width, height, overlap_percentage, num_inference_steps, resiz
         guidance_scale=3.5,
         negative_prompt="",
         true_guidance_scale=3.5,
-        output_type="latent"
-    ).images
-    print(latent_image)
-    pipe.to("cpu")
-    vae.to("cuda")
-    result = vae.decode(latent_image).sample
+    ).images[0]
+   
     result = result.convert("RGBA")
     cnet_image.paste(result, (0, 0), mask)