Lora loader

app.py

@@ -11,68 +11,72 @@ import spaces
 from comfy import model_management
 
 CHROMA_VERSION = "chroma-unlocked-v48-detail-calibrated.safetensors"
+LORA_NAME = "adapter_model.safetensors"
 
 # Download required models
 t5_path = hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp8_e4m3fn.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="lodestones/Chroma", filename=CHROMA_VERSION, local_dir="models/unet")
 
-# Example prompts with their parameters
+# Download LoRA file - you'll need to replace this with the correct repo_id and ensure the file exists
+lora_path = hf_hub_download(repo_id="zaksynack/chroma-kwis-v1", filename=LORA_NAME, local_dir="models/loras")
+
+# Example prompts with their parameters (updated to include LoRA strength)
 EXAMPLES = [
     [
         "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.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 36, 3.0, 229
+        1024, 1024, 36, 3.0, 229, 0.8
     ],
     [
         "A young man smiles broadly while wearing a blue captain's hat and blue mirrored aviator sunglasses. The man is the sharp focus of the image, set against a nicely blurred waterfront background featuring boats and buildings. The shot is cheerful, playful, and amateur looking.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 40, 2.1, 424
+        1024, 1024, 40, 2.1, 424, 0.7
     ],
     [
         "A Shiba Inu smiles broadly while wearing a blue captain's hat and blue mirrored aviator sunglasses. The dog is the sharp focus of the image, set against a nicely blurred waterfront background featuring boats and buildings. The shot is cheerful, playful, and amateur looking.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 28, 2.9, 413
+        1024, 1024, 28, 2.9, 413, 0.9
     ],
     [
         "A painting shows a bearded man in a beige shirt gripping a rope on a sailboat, battling choppy seas under a dramatic sunset. Another figure appears in the distant background, and the style is impressionistic with clear brushstrokes and high skill. The overall effect is one of adventure and freedom.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 33, 4.6, 281
+        1024, 1024, 33, 4.6, 281, 0.6
     ],
     [
         "A painting shows an anthro male wolf sitting on a rock in a forest, anthro, male, solo, canine, canid, werewolf, muscular anthro, black body, black fur, reaching at viewer with his paws, detailed background, outdoors, outside. The scene is detailed, with a shallow depth of field focusing on the anthro wolf. The overall effect is high fantasy and cinematic.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 26, 3.8, 326
+        1024, 1024, 26, 3.8, 326, 1.0
     ],
     [
         "A cute cartoon illustration of A massive black dragon with cosmic galaxy-patterned wings sits atop a castle overlooking a medieval port town at sunset. The style is bright and pastel-toned with soft lines. The artwork is clean and fairly well-executed.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 41, 5.7, 244
+        1024, 1024, 41, 5.7, 244, 0.5
     ],
     [
         "A soft longhaired male anthro dog, with a white chest patch and bright yellow-green eyes, sits on a light grey kitchen island in a sun-drenched NYC apartment. The city view is blurred, showcasing a shallow depth of field. The cat has a dark collar with a small bell. The photo's aesthetic is clean, modern, and slightly dramatic due to the lighting.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 26, 4.7, 242
+        1024, 1024, 26, 4.7, 242, 0.8
     ],
     [
         "A cheerful collage showcases 4 individual leopard geckos, each in a separate square. The geckos vary in color and pattern, but all are shown close-up and looking directly at the camera. Backgrounds differ, illustrating various terrarium elements and simple surfaces. High-quality shots, bright lighting, amateur photo quality. The overall aesthetic is heartwarming and simple.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 26, 4.7, 249
+        1024, 1024, 26, 4.7, 249, 0.7
     ],
     [
         "A cheerful collage showcases 4 cute anthro wolf nick wilde plushie, each in a separate square. Backgrounds differ, illustrating various bed and simple surfaces. High-quality shots, bright lighting, amateur photo quality. The overall aesthetic is heartwarming and simple.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 26, 5.4, 263
+        1024, 1024, 26, 5.4, 263, 0.9
     ],
     [
         "this picture depics a telegram sticker with thick white outlines of a male anthro fox",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 28, 3.9, 399
+        1024, 1024, 28, 3.9, 399, 0.8
     ],
     [
         "Extreme close-up of a single tiger eye, direct frontal view. Detailed iris and pupil. Sharp focus on eye texture and color. Natural lighting to capture authentic eye shine and depth. The word \"Chroma\" is painted over it in big, white brush strokes with visible texture.",
         "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
-        1024, 1024, 26, 4.0, 9
+        1024, 1024, 26, 4.0, 9, 0.6
     ]
 ]
 
@@ -154,7 +158,10 @@ samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
 vaedecode = VAEDecode()
 saveimage = SaveImage()
 
-# Load models
+# Initialize LoRA loader
+loraloader = NODE_CLASS_MAPPINGS["LoraLoader"]()
+
+# Load base models
 cliploader_78 = cliploader.load_clip(
     clip_name="t5xxl_fp8_e4m3fn.safetensors", type="chroma", device="default"
 )
@@ -180,32 +187,56 @@ valid_models = [
 model_management.load_models_gpu(valid_models)
 
 @spaces.GPU
-def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
+def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed, lora_strength):
     with torch.inference_mode():
         # Set random seed if provided
         if seed == -1:
             seed = random.randint(1, 2**64)
         random.seed(seed)
-        
+
+        # Load LoRA if strength > 0
+        if lora_strength > 0:
+            try:
+                loraloader_result = loraloader.load_lora(
+                    lora_name=LORA_NAME,
+                    strength_model=lora_strength,
+                    strength_clip=lora_strength,
+                    model=get_value_at_index(unetloader_76, 0),
+                    clip=get_value_at_index(t5tokenizeroptions_82, 0)
+                )
+                model_with_lora = get_value_at_index(loraloader_result, 0)
+                clip_with_lora = get_value_at_index(loraloader_result, 1)
+                print(f"LoRA {LORA_NAME} loaded with strength {lora_strength}")
+            except Exception as e:
+                print(f"Failed to load LoRA: {e}")
+                # Fall back to original model if LoRA loading fails
+                model_with_lora = get_value_at_index(unetloader_76, 0)
+                clip_with_lora = get_value_at_index(t5tokenizeroptions_82, 0)
+        else:
+            model_with_lora = get_value_at_index(unetloader_76, 0)
+            clip_with_lora = get_value_at_index(t5tokenizeroptions_82, 0)
+
         randomnoise_68 = randomnoise.get_noise(noise_seed=seed)
         emptysd3latentimage_69 = emptysd3latentimage.generate(
             width=width, height=height, batch_size=1
         )
         ksamplerselect_72 = ksamplerselect.get_sampler(sampler_name="euler")
 
+        # Use the clip with LoRA for text encoding
         cliptextencode_74 = cliptextencode.encode(
             text=prompt,
-            clip=get_value_at_index(t5tokenizeroptions_82, 0),
+            clip=clip_with_lora,
         )
 
         cliptextencode_75 = cliptextencode.encode(
             text=negative_prompt,
-            clip=get_value_at_index(t5tokenizeroptions_82, 0),
+            clip=clip_with_lora,
         )
 
+        # Use the model with LoRA for CFG guider
         cfgguider_73 = cfgguider.get_guider(
             cfg=cfg,
-            model=get_value_at_index(unetloader_76, 0),
+            model=model_with_lora,
             positive=get_value_at_index(cliptextencode_74, 0),
             negative=get_value_at_index(cliptextencode_75, 0),
         )
@@ -214,7 +245,7 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
             scheduler="beta",
             steps=steps,
             denoise=1,
-            model=get_value_at_index(unetloader_76, 0),
+            model=model_with_lora,
         )
 
         samplercustomadvanced_67 = samplercustomadvanced.sample(
@@ -232,10 +263,10 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
 
         # Save image using SaveImage node with simple string prefix
         saved = saveimage.save_images(
-            filename_prefix="Chroma_Generated",
+            filename_prefix="Chroma_LoRA_Generated",
             images=get_value_at_index(vaedecode_79, 0),
         )
-        
+
         # Return the path to the saved image
         saved_path = f"output/{saved['ui']['images'][0]['filename']}"
         return saved_path
@@ -243,17 +274,18 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
 # Create Gradio interface
 with gr.Blocks() as app:
     gr.Markdown(f"""
-# Chroma
+# Chroma with LoRA Support
 
 Model: [Chroma](https://huggingface.co/lodestones/Chroma) by [lodestones](https://huggingface.co/lodestones)
 
 Chroma Version: {CHROMA_VERSION}
+LoRA: {LORA_NAME}
 
 Run any ComfyUI Workflow on Spaces: [ComfyUI Workflows](https://huggingface.co/blog/run-comfyui-workflows-on-spaces)
 
 Space Author: [GitHub](https://github.com/gokayfem) | [X.com](https://x.com/gokayfem)
     """)
-    
+
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(
@@ -267,7 +299,7 @@ Space Author: [GitHub](https://github.com/gokayfem) | [X.com](https://x.com/goka
                 value="low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
                 lines=2
             )
-            
+
             with gr.Row():
                 width = gr.Slider(
                     minimum=512,
@@ -283,7 +315,7 @@ Space Author: [GitHub](https://github.com/gokayfem) | [X.com](https://x.com/goka
                     step=64,
                     label="Height"
                 )
-            
+
             with gr.Row():
                 steps = gr.Slider(
                     minimum=1,
@@ -299,30 +331,39 @@ Space Author: [GitHub](https://github.com/gokayfem) | [X.com](https://x.com/goka
                     step=0.5,
                     label="CFG Scale"
                 )
+
+            with gr.Row():
                 seed = gr.Number(
                     value=-1,
                     label="Seed (-1 for random)"
                 )
-            
+                lora_strength = gr.Slider(
+                    minimum=0,
+                    maximum=2.0,
+                    value=0.8,
+                    step=0.1,
+                    label="LoRA Strength (0 = disabled)"
+                )
+
             generate_btn = gr.Button("Generate")
-        
+
         with gr.Column():
             output_image = gr.Image(label="Generated Image")
-    
+
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed],
+        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed, lora_strength],
         outputs=[output_image]
     )
-    
+
     # Add examples section
     gr.Examples(
         examples=EXAMPLES,
-        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed],
+        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed, lora_strength],
         outputs=[output_image],
         fn=generate_image,
         cache_examples=True,
-        label="Example Prompts - Click to try!"
+        label="Example Prompts with LoRA - Click to try!"
     )
 
 if __name__ == "__main__":