Update comfy_pulid.py

comfy_pulid.py

@@ -3,6 +3,17 @@ import random
 import sys
 from typing import Sequence, Mapping, Any, Union
 import torch
+from comfy import model_management
+from huggingface_hub import hf_hub_download
+import spaces
+
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-Redux-dev", filename="flux1-redux-dev.safetensors", local_dir="models/style_models")
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-Depth-dev", filename="flux1-depth-dev.safetensors", local_dir="models/diffusion_models")
+hf_hub_download(repo_id="Comfy-Org/sigclip_vision_384", filename="sigclip_vision_patch14_384.safetensors", local_dir="models/clip_vision")
+hf_hub_download(repo_id="Kijai/DepthAnythingV2-safetensors", filename="depth_anything_v2_vitl_fp32.safetensors", local_dir="models/depthanything")
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="ae.safetensors", local_dir="models/vae/FLUX1")
+hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.safetensors", local_dir="models/text_encoders")
+hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir="models/text_encoders/t5")
 
 
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
@@ -114,12 +125,94 @@ def import_custom_nodes() -> None:
 
 from nodes import NODE_CLASS_MAPPINGS
 
+intconstant = NODE_CLASS_MAPPINGS["INTConstant"]()
+dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+
+#To be added to `model_loaders` as it loads a model
+dualcliploader_357 = dualcliploader.load_clip(
+    clip_name1="t5/t5xxl_fp16.safetensors",
+    clip_name2="clip_l.safetensors",
+    type="flux",
+)
+cr_clip_input_switch = NODE_CLASS_MAPPINGS["CR Clip Input Switch"]()
+cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+imageresize = NODE_CLASS_MAPPINGS["ImageResize+"]()
+getimagesizeandcount = NODE_CLASS_MAPPINGS["GetImageSizeAndCount"]()
+vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+
+#To be added to `model_loaders` as it loads a model
+vaeloader_359 = vaeloader.load_vae(vae_name="FLUX1/ae.safetensors")
+
+vaeencode = NODE_CLASS_MAPPINGS["VAEEncode"]()
+unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+
+#To be added to `model_loaders` as it loads a model
+unetloader_358 = unetloader.load_unet(
+    unet_name="flux1-depth-dev.safetensors", weight_dtype="default"
+)
+ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
+randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
+fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
+depthanything_v2 = NODE_CLASS_MAPPINGS["DepthAnything_V2"]()
+downloadandloaddepthanythingv2model = NODE_CLASS_MAPPINGS[
+    "DownloadAndLoadDepthAnythingV2Model"
+]()
+
+#To be added to `model_loaders` as it loads a model
+downloadandloaddepthanythingv2model_437 = (
+    downloadandloaddepthanythingv2model.loadmodel(
+        model="depth_anything_v2_vitl_fp32.safetensors"
+    )
+)
+instructpixtopixconditioning = NODE_CLASS_MAPPINGS[
+    "InstructPixToPixConditioning"
+]()
+text_multiline_454 = text_multiline.text_multiline(text="FLUX_Redux")
+clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+
+#To be added to `model_loaders` as it loads a model
+clipvisionloader_438 = clipvisionloader.load_clip(
+    clip_name="sigclip_vision_patch14_384.safetensors"
+)
+clipvisionencode = NODE_CLASS_MAPPINGS["CLIPVisionEncode"]()
+stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+
+#To be added to `model_loaders` as it loads a model
+stylemodelloader_441 = stylemodelloader.load_style_model(
+    style_model_name="flux1-redux-dev.safetensors"
+)
+text_multiline = NODE_CLASS_MAPPINGS["Text Multiline"]()
+emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
+cr_conditioning_input_switch = NODE_CLASS_MAPPINGS[
+    "CR Conditioning Input Switch"
+]()
+cr_model_input_switch = NODE_CLASS_MAPPINGS["CR Model Input Switch"]()
+stylemodelapplyadvanced = NODE_CLASS_MAPPINGS["StyleModelApplyAdvanced"]()
+basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]()
+basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
+samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
+vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
+imagecrop = NODE_CLASS_MAPPINGS["ImageCrop+"]()
+
+#Add all the models that load a safetensors file
+model_loaders = [dualcliploader_357, vaeloader_359, unetloader_358, clipvisionloader_438, stylemodelloader_441, downloadandloaddepthanythingv2model_437]
+
+# Check which models are valid and how to best load them
+valid_models = [
+    getattr(loader[0], 'patcher', loader[0]) 
+    for loader in model_loaders
+    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
+]
+
+#Finally loads the models
+model_management.load_models_gpu(valid_models)
+
 
 def generate_image(prompt, structure_image, style_image, depth_strength, style_strength):
     import_custom_nodes()
     with torch.inference_mode():
-        vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-        vaeloader_10 = vaeloader.load_vae(vae_name="FLUX1/ae.safetensors")
 
         dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
         dualcliploader_11 = dualcliploader.load_clip(
@@ -284,5 +377,38 @@ def generate_image(prompt, structure_image, style_image, depth_strength, style_s
         return saved_path
 
 
-#if __name__ == "__main__":
-#    main()
+if __name__ == "__main__":
+    with gr.Blocks() as app:
+            # Add a title
+            gr.Markdown("# FLUX Style Shaping")
+    
+            with gr.Row():
+                with gr.Column():
+                    # Add an input
+                    prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
+                    # Add a `Row` to include the groups side by side 
+                    with gr.Row():
+                        # First group includes structure image and depth strength
+                        with gr.Group():
+                            structure_image = gr.Image(label="Structure Image", type="filepath")
+                            depth_strength = gr.Slider(minimum=0, maximum=50, value=15, label="Depth Strength")
+                        # Second group includes style image and style strength
+                        with gr.Group():
+                            style_image = gr.Image(label="Style Image", type="filepath")
+                            style_strength = gr.Slider(minimum=0, maximum=1, value=0.5, label="Style Strength")
+    
+                    # The generate button
+                    generate_btn = gr.Button("Generate")
+    
+                with gr.Column():
+                    # The output image
+                    output_image = gr.Image(label="Generated Image")
+    
+                # When clicking the button, it will trigger the `generate_image` function, with the respective inputs
+                # and the output an image
+                generate_btn.click(
+                    fn=generate_image,
+                    inputs=[prompt_input, structure_image, style_image, depth_strength, style_strength],
+                    outputs=[output_image]
+                )
+            app.launch(share=True)