diff flux.py @ 0:0d0561746128 draft

planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools/flux commit 71b3dacee16dc999cb4fa113858d6ace1781c71c

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/flux.py	Mon Oct 14 16:50:38 2024 +0000
@@ -0,0 +1,32 @@
+import sys
+
+import torch
+from diffusers import FluxPipeline
+
+model = sys.argv[1]
+
+prompt_type = sys.argv[2]
+if prompt_type == "file":
+    with open(sys.argv[3], "r") as f:
+        prompt = f.read().strip()
+elif prompt_type == "text":
+    prompt = sys.argv[3]
+
+if model not in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
+    print("Invalid model!")
+    sys.exit(1)
+
+
+pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
+pipe.to(torch.float16)
+
+image = pipe(
+    prompt,
+    num_inference_steps=4,
+    generator=torch.Generator("cpu").manual_seed(42),
+).images[0]
+
+image.save("output.png")