Update raw.py

raw.py

@@ -2,19 +2,22 @@ import torch
 import spaces
 import os
 from diffusers.utils import load_image
-from diffusers import FluxControlNetModel, FluxControlNetPipeline
+from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 import gradio as gr
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to(device)
+
 # Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler",
     torch_dtype=torch.bfloat16
 )
 pipe = FluxControlNetPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
+    "LPX55/FLUX.1-merged_uncensored",
     controlnet=controlnet,
     torch_dtype=torch.bfloat16,
+    vae=good_vae,
     token=huggingface_token
 )
 pipe.to("cuda")
@@ -25,7 +28,7 @@ def generate_image(prompt, scale, steps, control_image):
     control_image = load_image(control_image)
     w, h = control_image.size
     # Upscale x1
-    control_image = control_image.resize((w * scale, h * scale))
+    control_image = control_image.resize((int(w * scale), int(h * scale)))
     image = pipe(
         prompt=prompt,
         control_image=control_image,