Update README.md

README.md

@@ -53,6 +53,73 @@ image = pipe(
 ).images[0]
 image.save("chroma.png")
 ```
+
+Quantized inference using gemlite
+
+```py
+import torch
+from diffusers import ChromaPipeline
+
+pipe = ChromaPipeline.from_pretrained("lodestones/Chroma1-HD", torch_dtype=torch.float16)
+#pipe.enable_model_cpu_offload()
+
+#######################################################
+import gemlite
+device = 'cuda:0'
+processor = gemlite.helper.A8W8_int8_dynamic
+#processor = gemlite.helper.A8W8_fp8_dynamic
+#processor = gemlite.helper.A16W4_MXFP
+
+for name, module in pipe.transformer.named_modules():
+    module.name = name
+
+def patch_linearlayers(model, fct):
+    for name, layer in model.named_children():
+        if isinstance(layer, torch.nn.Linear):
+            setattr(model, name, fct(layer, name))
+        else:
+            patch_linearlayers(layer, fct)
+
+def patch_linear_to_gemlite(layer, name):
+    layer = layer.to(device, non_blocking=True)
+    try:
+        return processor(device=device).from_linear(layer) 
+    except Exception as exception:
+        print('Skipping gemlite conversion for: ' + str(layer.name), exception)
+        return layer
+
+patch_linearlayers(pipe.transformer, patch_linear_to_gemlite)
+torch.cuda.synchronize()
+torch.cuda.empty_cache()
+
+pipe.to(device)
+pipe.transformer.forward = torch.compile(pipe.transformer.forward, fullgraph=True)
+pipe.vae.forward = torch.compile(pipe.vae.forward, fullgraph=True)
+#pipe.set_progress_bar_config(disable=True)
+#######################################################
+
+prompt = [
+    "A high-fashion close-up portrait of a blonde woman in clear sunglasses. The image uses a bold teal and red color split for dramatic lighting. The background is a simple teal-green. The photo is sharp and well-composed, and is designed for viewing with anaglyph 3D glasses for optimal effect. It looks professionally done."
+]
+negative_prompt =  ["low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors"]
+
+import time
+for _ in range(3):
+    t_start = time.time()
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        generator=torch.Generator("cpu").manual_seed(433),
+        num_inference_steps=40,
+        guidance_scale=3.0,
+        num_images_per_prompt=1,
+    ).images[0]
+    t_end = time.time()
+    print(f"Took: {t_end - t_start} secs.") #66.1242527961731 -> 27.72 sec
+
+image.save("chroma.png")
+```
+
 ComfyUI
 For advanced users and customized workflows, you can use Chroma with ComfyUI.