change to vae

app.py

@@ -2,19 +2,15 @@ import gradio as gr
 import numpy as np
 
 import spaces
-import torch
 import random
 from image_gen_aux import DepthPreprocessor
+from PIL import Image
+import torch
+from torchvision import transforms
 
-from diffusers import FluxFillPipeline
+from diffusers import FluxFillPipeline, AutoencoderKL
 from PIL import Image
 
-def remove_background(image):
-       # Placeholder function for background removal
-       # Use a library or model like Inspyrenet for actual implementation
-       mask = generate_mask(image)
-       subject = apply_mask(image, mask)
-       return subject, mask
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -23,8 +19,29 @@ pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", tor
 pipe.load_lora_weights("alvdansen/flux-koda")
 pipe.enable_lora()
 
+vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae")
 processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
 
+preprocess = transforms.Compose(
+    [
+        transforms.Resize(
+            (vae.config.sample_size, vae.config.sample_size),
+            interpolation=transforms.InterpolationMode.BILINEAR,
+        ),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
+#
+# image_np = image[0].cpu().numpy()  # Move to CPU and convert to NumPy
+
+# if image_np.shape[0] == 3:  # Check if channels are first
+#     image_np = image_np.transpose(1, 2, 0)
+
+# image_np = (image_np * 255).astype(np.uint8)
+
+image = Image.fromarray(image_np)
+
 def calculate_optimal_dimensions(image: Image.Image):
     # Extract the original dimensions
     original_width, original_height = image.size
@@ -74,7 +91,7 @@ def infer(edit_images, prompt, seed=42, randomize_seed=False, width=1024, height
 
     controlImage = processor(image)[0].convert("RGB")
     image = pipe(
-        control_image=controlImage,
+        mask_image_latent=vae.encode(controlImage),
         prompt=prompt,
         image=image,
         mask_image=mask,

requirements.txt

@@ -6,3 +6,5 @@ safetensors
 sentencepiece
 peft
 xformers
+torchvision
+torch