Refactor flux.py and flux.xml for improved model selection (bgruening…

…#1536)

tools/flux/flux.py

@@ -1,9 +1,10 @@
+import os
 import sys
 
 import torch
 from diffusers import FluxPipeline
 
-model = sys.argv[1]
+model_path = sys.argv[1]
 
 prompt_type = sys.argv[2]
 if prompt_type == "file":
@@ -12,20 +13,28 @@
 elif prompt_type == "text":
     prompt = sys.argv[3]
 
-if model not in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
+if "dev" in model_path:
+    num_inference_steps = 20
+elif "schnell" in model_path:
+    num_inference_steps = 4
+else:
     print("Invalid model!")
     sys.exit(1)
 
+snapshots = []
+for d in os.listdir(os.path.join(model_path, "snapshots")):
+    snapshots.append(os.path.join(model_path, "snapshots", d))
+latest_snapshot_path = max(snapshots, key=os.path.getmtime)
 
-pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+pipe = FluxPipeline.from_pretrained(latest_snapshot_path, 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,
+    num_inference_steps=num_inference_steps,
     generator=torch.Generator("cpu").manual_seed(42),
 ).images[0]
 

tools/flux/flux.xml

@@ -2,7 +2,7 @@
     <description>text-to-image model</description>
     <macros>
         <token name="@TOOL_VERSION@">2024</token>
-        <token name="@VERSION_SUFFIX@">0</token>
+        <token name="@VERSION_SUFFIX@">1</token>
     </macros>
     <requirements>
         <requirement type="package" version="3.12">python</requirement>
@@ -16,9 +16,8 @@
         <requirement type="package" version="0.24.6">huggingface_hub</requirement>
     </requirements>
     <command detect_errors="exit_code"><![CDATA[
-export HF_HOME='$flux_models.fields.path' &&
 python '$__tool_directory__/flux.py'
-'$flux_models'
+'$flux_models.fields.path'
 '$input_type_selector'
 '$prompt'
     ]]></command>