Update app.py

app.py

@@ -11,11 +11,12 @@ import spaces
 from comfy import model_management
 
 CHROMA_VERSION = "Chroma1-HD-flash-heun/Chroma1-HD-flash-heun-fp8_scaled_original_hybrid_large.safetensors"
-
+LORA_VERSION = "flash-heun-pruned/chroma-flash-heun_r01-fp32-pruned.safetensors"
 # Download required models
 t5_path = hf_hub_download(repo_id="silveroxides/flan-t5-xxl-encoder-only", filename="flan-t5-xxl_float8_e4m3fn_scaled_stochastic.safetensors", local_dir="models/text_encoders/")
 vae_path = hf_hub_download(repo_id="lodestones/Chroma", filename="ae.safetensors", local_dir="models/vae")
 unet_path = hf_hub_download(repo_id="silveroxides/Chroma-Misc-Models", filename=CHROMA_VERSION, local_dir="models/unet")
+lora_path = hf_hub_download(repo_id="silveroxides/Chroma-LoRAs", filename=LORA_VERSION, local_dir="models/loras")
 
 # Example prompts with their parameters
 EXAMPLES = [
@@ -134,7 +135,7 @@ from nodes import (
     CLIPTextEncode,
     CLIPLoader,
     VAEDecode,
-    UNETLoader,
+    LoraLoaderModelOnly,
     VAELoader,
     SaveImage,
 )
@@ -148,6 +149,7 @@ t5tokenizeroptions = NODE_CLASS_MAPPINGS["T5TokenizerOptions"]()
 cliptextencode = CLIPTextEncode()
 unetloader = NODE_CLASS_MAPPINGS["ScaledFP8HybridUNetLoader"]()
 vaeloader = VAELoader()
+loraloader = LoraLoaderModelOnly()
 modelsamplingauraflow = NODE_CLASS_MAPPINGS["ModelSamplingAuraFlow"]()
 cfgguider = NODE_CLASS_MAPPINGS["CFGGuider"]()
 betascheduler = NODE_CLASS_MAPPINGS["BetaSamplingScheduler"]()
@@ -160,15 +162,24 @@ cliploader_78 = cliploader.load_clip(
     clip_name="flan-t5-xxl_float8_e4m3fn_scaled_stochastic.safetensors", type="chroma", device="default"
 )
 t5tokenizeroptions_82 = t5tokenizeroptions.set_options(
-    min_padding=0, min_length=0, clip=get_value_at_index(cliploader_78, 0)
+    min_padding=0,
+    min_length=0,
+    clip=get_value_at_index(cliploader_78, 0)
 )
 unetloader_76 = unetloader.load_unet(
     model_name=CHROMA_VERSION, chroma_hybrid_large=True, radiance_hybrid_large=False, chroma_hybrid_small=False, radiance_hybrid_small=False, wan=False, qwen=False
 )
 vaeloader_80 = vaeloader.load_vae(vae_name="ae.safetensors")
 
+loraloader_81 = loraloader.load_lora_model_only(
+    lora_name=LORA_VERSION,
+    strength_model=1.0,
+    model=get_value_at_index(unetloader_76, 0),
+)
+
+
 # Add all the models that load a safetensors file
-model_loaders = [cliploader_78, unetloader_76, vaeloader_80]
+model_loaders = [cliploader_78, unetloader_76, vaeloader_80, loraloader_81]
 
 # Check which models are valid and how to best load them
 valid_models = [
@@ -203,15 +214,14 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
             clip=get_value_at_index(t5tokenizeroptions_82, 0),
         )
 
-
         modelsamplingauraflow_85 = modelsamplingauraflow.patch(
             shift=1.0,
-            model=get_value_at_index(unetloader_76, 0),
+            model=get_value_at_index(loraloader_81, 0),
         )
 
         cfgguider_73 = cfgguider.get_guider(
             cfg=cfg,
-            model=get_value_at_index(unetloader_76, 0),
+            model=get_value_at_index(loraloader_81, 0),
             positive=get_value_at_index(cliptextencode_74, 0),
             negative=get_value_at_index(cliptextencode_75, 0),
         )
@@ -220,7 +230,7 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
             steps=steps,
             alpha=0.45,
             beta=0.45,
-            model=get_value_at_index(unetloader_76, 0),
+            model=get_value_at_index(loraloader_81, 0),
         )
 
         samplercustomadvanced_67 = samplercustomadvanced.sample(