Update app.py

app.py

@@ -28,9 +28,7 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
 MAX_SEED = 2**32-1
 
@@ -88,10 +86,10 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
             width=width,
             height=height,
             generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-            output_type="pil",
-            good_vae=good_vae,
-        )
+            joint_attention_kwargs={"scale": lora_scale}
+        ).images[0]
+        
+    progress(99, "Generate success!")
     return generate_image
 
 
@@ -117,6 +115,8 @@ def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, see
     else:
         result = {"status": "success", "message": "Image generated but not uploaded"}
     
+    progress(100, "Completed!")
+
     yield final_image, seed, gr.update(value=progress_bar, visible=False),  json.dumps(result)