Update app.py

app.py

@@ -26,15 +26,22 @@ clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
 import numpy as np
 import random
 import torch
-from diffusers import FluxPipeline
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload()
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
+torch.cuda.empty_cache()
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-#pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 @spaces.GPU(duration=75)
 def infer_flux(prompt, seed=42, randomize_seed=True, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
@@ -42,18 +49,18 @@ def infer_flux(prompt, seed=42, randomize_seed=True, width=1024, height=1024, gu
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
     
-    image = pipe(
-        prompt,
-        height=1024,
-        width=1024,
-        guidance_scale=3.5,
-        num_inference_steps=25,
-        max_sequence_length=512,
-        generator=torch.Generator("cpu").manual_seed(0)
-    ).images[0]
-    image.save("flux-dev.png")
-    return "flux-dev.png"
-
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            output_type="pil",
+            good_vae=good_vae,
+        ):
+            yield img
+    
 @spaces.GPU
 def llama_gen_fragrance(scene):
 
@@ -281,14 +288,6 @@ def infer(image_input):
     image_desc = extract_field(parsed, "Image Description")
     yield result, parsed, image_desc
 
-def generate_flacon(image_desc, progress=gr.Progress(track_tqdm=True)):
-    
-    print(image_desc)
-    gr.Info("Generating a nice marketing image with FLUX...")
-    gen_bottle = infer_flux(image_desc)
-
-    return gen_bottle
-
 css="""
 #col-container {max-width: 910px; margin-left: auto; margin-right: auto;}
 """
@@ -314,6 +313,6 @@ with gr.Blocks(css=css) as demo:
                 bottle_res = gr.Image(label="Flacon")
         
     submit_btn.click(fn=infer, inputs=[image_in], outputs=[fragrance, json_res, flacon_desc])
-    get_flacon_btn.click(fn=generate_flacon, inputs=[flacon_desc], outputs=[bottle_res])
+    get_flacon_btn.click(fn=infer_flux, inputs=[flacon_desc], outputs=[bottle_res])
 
 demo.queue(max_size=12).launch(ssr_mode=False, mcp_server=True)