Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import os
 
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import AutoPipelineForText2Image
+from diffusers import AutoPipelineForText2Image, AutoencoderKL, AutoencoderTiny
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -12,11 +12,14 @@ hf_token = os.getenv('HF_TOKEN')
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
+    torch.cuda.empty_cache()
 else:
     torch_dtype = torch.float32
 
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 pipe = pipeline = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.1-dev", token=hf_token, torch_dtype=torch.bfloat16)
 pipe.load_lora_weights('aleksa-codes/flux-ghibsky-illustration', weight_name='lora.safetensors')
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -26,13 +29,12 @@ MAX_IMAGE_SIZE = 1024
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
+    seed=42,
+    randomize_seed=True,
+    width=1024,
+    height=1024,
+    guidance_scale=3.5,
+    num_inference_steps=28,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -42,12 +44,12 @@ def infer(
 
     image = pipe(
         prompt=prompt,
-        negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,
         height=height,
         generator=generator,
+        good_vae=good_vae,
     ).images[0]
 
     return image, seed