update gradio & add sd_embed

README.md

@@ -4,7 +4,7 @@ emoji: 🦢
 colorFrom: gray
 colorTo: gray
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -6,14 +6,26 @@ import gradio as gr
 import numpy as np
 import random
 import torch
-from diffusers import FluxPipeline
+from diffusers import FluxPipeline, DiffusionPipeline, FluxTransformer2DModel  # noqa: F401
+from torchao.quantization import quantize_, int8_weight_only
+from sd_embed.embedding_funcs import get_weighted_text_embeddings_flux1
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    torch_dtype=dtype,
+
+transformer = FluxTransformer2DModel.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", subfolder="transformer", torch_dtype=torch.bfloat16
+)
+quantize_(transformer, int8_weight_only())
+
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
 )
+
+# pipe = FluxPipeline.from_pretrained(
+#     "black-forest-labs/FLUX.1-dev",
+#     torch_dtype=dtype,
+# )
 pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -35,8 +47,14 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
+
+    prompt_embeds, pooled_prompt_embeds = get_weighted_text_embeddings_flux1(
+        pipe=pipe, prompt=prompt
+    )
+
     image = pipe(
-        prompt=prompt,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
         width=width,
         height=height,
         num_inference_steps=num_inference_steps,

requirements.txt

@@ -1,7 +1,10 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
 accelerate
 git+https://github.com/huggingface/diffusers.git
 torch
-transformers==4.42.4
+transformers
 xformers
 sentencepiece
-bitsandbytes
+bitsandbytes
+git+https://github.com/xhinker/sd_embed.git@main
+torchao