Add Inference Example Scripts

inference/example_full.py

@@ -0,0 +1,31 @@
+# Modified from https://huggingface.co/black-forest-labs/FLUX.1-dev usage example
+import torch
+from diffusers import FluxPipeline, FluxTransformer2DModel
+
+# 1. Load pipeline
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16
+)
+
+# 2. Load full fine-tuned transformer weights
+checkpoint = 'path/to/checkpoint'  # replace with your checkpoint directory
+pipe.transformer = FluxTransformer2DModel.from_pretrained(
+    checkpoint,
+    torch_dtype=torch.bfloat16
+)
+
+pipe.enable_model_cpu_offload()
+
+# 3. Generate image
+prompt = "A cat holding a sign that says hello world"
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    guidance_scale=3.5,
+    num_inference_steps=28,
+    max_sequence_length=512,
+    generator=torch.Generator("cpu").manual_seed(0)
+).images[0]
+image.save("flux-dev-full.png")

inference/example_lora.py

@@ -0,0 +1,33 @@
+# Modified from https://huggingface.co/black-forest-labs/FLUX.1-dev usage example
+import torch
+from diffusers import FluxPipeline
+from peft import PeftModel
+
+# 1. Load pipeline
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16
+)
+
+# 2. Load LoRA weights
+checkpoint = 'path/to/checkpoint'  # replace with your checkpoint directory
+pipe.transformer = PeftModel.from_pretrained(
+    pipe.transformer,
+    checkpoint,
+    torch_dtype=torch.bfloat16
+)
+
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+
+# 3. Generate image
+prompt = "A cat holding a sign that says hello world"
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    guidance_scale=3.5,
+    num_inference_steps=28,
+    max_sequence_length=512,
+    generator=torch.Generator("cpu").manual_seed(0)
+).images[0]
+image.save("flux-dev.png")