Update app.py

app.py

@@ -21,6 +21,46 @@ model = AutoModelForCausalLM.from_pretrained(model_path, use_auth_token=hf_token
 
 clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")
 
+# FLUX
+
+import numpy as np
+import random
+import torch
+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
+
+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)
+
+@spaces.GPU
+def infer_flux(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    
+    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):
 
@@ -247,8 +287,9 @@ def infer(image_input):
     image_desc = extract_field(parsed, "Image Description")
 
     print(image_desc)
+    gen_bottle = infer_flux(image_desc)
 
-    return result, parsed
+    return result, parsed, gen_bottle
 
 css="""
 #col-container {max-width: 910px; margin-left: auto; margin-right: auto;}
@@ -271,7 +312,8 @@ with gr.Blocks(css=css) as demo:
                 #caption = gr.Textbox(label="Generated Caption")
                 fragrance = gr.Textbox(label="generated Fragrance", elem_id="fragrance")
                 json_res = gr.JSON(label="JSON")
+                bottle_res = gr.Image(label="Flacon")
         
-    submit_btn.click(fn=infer, inputs=[image_in], outputs=[fragrance, json_res])
+    submit_btn.click(fn=infer, inputs=[image_in], outputs=[fragrance, json_res, bottle_res])
 
 demo.queue(max_size=12).launch(ssr_mode=False, mcp_server=True)