Commit

app.py

@@ -1,7 +1,163 @@
+import spaces
 import gradio as gr
+import numpy as np
+import random
+import python
+import spaces
+import torch
+import os
+from huggingface_hub import hf_hub_download
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
+from peft import PeftModel
 
-def greet(name):
-    return "Hello " + name + "!!"
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+token = os.getenv("HUGGINGFACE_TOKEN")
+
+# good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype, token=token).to(device)
+# pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token=token).to(device)
+torch.cuda.empty_cache()
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048  # not used anymore
+
+# Bind the custom method
+# pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+python.model_loading()
+
+
+@spaces.GPU()
+def infer(prompt, seed=42, randomize_seed=False, aspect_ratio="4:3 landscape 1152x896", lora_weight="lora_weight_rank_32_alpha_32.safetensors",
+          guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+    # Default height + width
+    width, height = 1024, 1024
+    
+    # Randomize seed if requested
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    
+    # Load the selected LoRA weight and fuse it
+    lora_weight_path = os.path.join("lora_weights", lora_weight)
+    # pipe.load_lora_weights(weight_path)
+    # pipe.fuse_lora()
+    torch.cuda.empty_cache()
+    image = python.generate_image(
+        prompt,
+        height,
+        width,
+        aspect_ratio,
+        seed,
+        guidance_scale=guidance_scale,
+    ).images[0]
+    # Generate images
+    # for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+    #         prompt=prompt,
+    #         guidance_scale=guidance_scale,
+    #         num_inference_steps=num_inference_steps,
+    #         width=width,
+    #         height=height,
+    #         generator=generator,
+    #         output_type="pil",
+    #         good_vae=good_vae,
+    #     ):
+    #     out_img = img
+    return image,seed
+
+# Examples for the prompt
+examples = [
+    "Photo on a small glass panel. Color. A vintage Autochrome photograph, early 1900s aesthetic depicts four roses in a brown vase with dark background.",
+    "Colorized photograph on a small glass panel depicting trees with orange leaves, a dirt path, and a wood and rope fence.",
+]
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# Autochrome image generator demo using FLUX.1 [dev]  
+[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
+        """)
+        
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=5,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Run", scale=0)
+        
+        result = gr.Image(label="Result", show_label=False)
+        
+        with gr.Accordion("Advanced Settings", open=False):
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            
+            # Dropdown for aspect ratio selection
+            aspect_ratio = gr.Dropdown(
+                label="Aspect Ratio",
+                choices=["1:1 square 1024x1024", "3:4 portrait 896x1152", "5:8 portrait 832x1216", "9:16 portrait 768x1344", "4:3 landscape 1152x896", "3:2 landscape 1216x832", "16:9 landscape 1344x768"],
+                value="4:3 landscape 1152x896",
+                interactive=True,
+            )
+            
+            # Dropdown for LoRA weight selection
+            lora_weight = gr.Dropdown(
+                label="LoRA Weight",
+                choices=[
+                    "lora_weight_rank_16_alpha_32_1.safetensors",
+                    "lora_weight_rank_16_alpha_32_2.safetensors",
+                    "lora_weight_rank_32_alpha_32.safetensors",
+                    "lora_weight_rank_32_alpha_64.safetensors",
+                ],
+                value="lora_weight_rank_32_alpha_32.safetensors",
+                interactive=True,
+            )
+            
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=25,
+                    step=0.1,
+                    value=6,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=40,
+                )
+        
+        gr.Examples(
+            examples=examples,
+            fn=infer,
+            inputs=[prompt],
+            outputs=[result, seed],
+            cache_examples="lazy"
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, seed, randomize_seed, aspect_ratio, lora_weight, guidance_scale, num_inference_steps],
+        outputs=[result, seed]
+        
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

custom_nodes/ComfyUI-to-Python-Extension

@@ -0,0 +1 @@
+Subproject commit 0aa2747736193939a3e1e8ef35aa3d0e378c60db

custom_nodes/ComfyUI_Comfyroll_CustomNodes

@@ -0,0 +1 @@
+Subproject commit d78b780ae43fcf8c6b7c6505e6ffb4584281ceca

live_preview_helpers.py

@@ -0,0 +1,166 @@
+import torch
+import numpy as np
+from diffusers import FluxPipeline, AutoencoderTiny, FlowMatchEulerDiscreteScheduler
+from typing import Any, Dict, List, Optional, Union
+
+# Helper functions
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+# FLUX pipeline function
+@torch.inference_mode()
+def flux_pipe_call_that_returns_an_iterable_of_images(
+    self,
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 28,
+    timesteps: List[int] = None,
+    guidance_scale: float = 3.5,
+    num_images_per_prompt: Optional[int] = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.FloatTensor] = None,
+    prompt_embeds: Optional[torch.FloatTensor] = None,
+    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    max_sequence_length: int = 512,
+    good_vae: Optional[Any] = None,
+):
+    height = height or self.default_sample_size * self.vae_scale_factor
+    width = width or self.default_sample_size * self.vae_scale_factor
+
+    # 1. Check inputs
+    self.check_inputs(
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        max_sequence_length=max_sequence_length,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._joint_attention_kwargs = joint_attention_kwargs
+    self._interrupt = False
+
+    # 2. Define call parameters
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    device = self._execution_device
+
+    # 3. Encode prompt
+    lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
+    prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
+        prompt=prompt,
+        prompt_2=prompt_2,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        max_sequence_length=max_sequence_length,
+        lora_scale=lora_scale,
+    )
+    # 4. Prepare latent variables
+    num_channels_latents = self.transformer.config.in_channels // 4
+    latents, latent_image_ids = self.prepare_latents(
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    image_seq_len = latents.shape[1]
+    mu = calculate_shift(
+        image_seq_len,
+        self.scheduler.config.base_image_seq_len,
+        self.scheduler.config.max_image_seq_len,
+        self.scheduler.config.base_shift,
+        self.scheduler.config.max_shift,
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        self.scheduler,
+        num_inference_steps,
+        device,
+        timesteps,
+        sigmas,
+        mu=mu,
+    )
+    self._num_timesteps = len(timesteps)
+
+    # Handle guidance
+    guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
+
+    # 6. Denoising loop
+    for i, t in enumerate(timesteps):
+        if self.interrupt:
+            continue
+
+        timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+        noise_pred = self.transformer(
+            hidden_states=latents,
+            timestep=timestep / 1000,
+            guidance=guidance,
+            pooled_projections=pooled_prompt_embeds,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=text_ids,
+            img_ids=latent_image_ids,
+            joint_attention_kwargs=self.joint_attention_kwargs,
+            return_dict=False,
+        )[0]
+        # Yield intermediate result
+        latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents_for_image, return_dict=False)[0]
+        yield self.image_processor.postprocess(image, output_type=output_type)[0]
+        
+        latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        torch.cuda.empty_cache()
+
+    # Final image using good_vae
+    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+    latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
+    image = good_vae.decode(latents, return_dict=False)[0]
+    self.maybe_free_model_hooks()
+    torch.cuda.empty_cache()
+    yield self.image_processor.postprocess(image, output_type=output_type)[0]

python.py

@@ -0,0 +1,262 @@
+import os
+import random
+import sys
+from typing import Sequence, Mapping, Any, Union
+import torch
+import spaces
+# from comfy import model_management
+from nodes import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_1
+from comfy_extras.nodes_custom_sampler import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_2
+from custom_nodes.ComfyUI_Comfyroll_CustomNodes.node_mappings import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_3
+from huggingface_hub import hf_hub_download
+
+# Merge both mappings
+COMBINED_NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS_1, **NODE_CLASS_MAPPINGS_2, **NODE_CLASS_MAPPINGS_3}
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="flux1-dev.safetensors", local_dir="models/unet")
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="ae.safetensors", local_dir="models/vae")
+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")
+
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    """Returns the value at the given index of a sequence or mapping.
+
+    If the object is a sequence (like list or string), returns the value at the given index.
+    If the object is a mapping (like a dictionary), returns the value at the index-th key.
+
+    Some return a dictionary, in these cases, we look for the "results" key
+
+    Args:
+        obj (Union[Sequence, Mapping]): The object to retrieve the value from.
+        index (int): The index of the value to retrieve.
+
+    Returns:
+        Any: The value at the given index.
+
+    Raises:
+        IndexError: If the index is o of bounds for the object and the object is not a mapping.
+    """
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+
+def find_path(name: str, path: str = None) -> str:
+    """
+    Recursively looks at parent folders starting from the given path until it finds the given name.
+    Returns the path as a Path object if found, or None otherwise.
+    """
+    # If no path is given, use the current working directory
+    if path is None:
+        path = os.getcwd()
+
+    # Check if the current directory contains the name
+    if name in os.listdir(path):
+        path_name = os.path.join(path, name)
+        print(f"{name} found: {path_name}")
+        return path_name
+
+    # Get the parent directory
+    parent_directory = os.path.dirname(path)
+
+    # If the parent directory is the same as the current directory, we've reached the root and stop the search
+    if parent_directory == path:
+        return None
+
+    # Recursively call the function with the parent directory
+    return find_path(name, parent_directory)
+
+
+def add_comfyui_directory_to_sys_path() -> None:
+    """
+    Add 'ComfyUI' to the sys.path
+    """
+    comfyui_path = find_path("ComfyUI")
+    if comfyui_path is not None and os.path.isdir(comfyui_path):
+        sys.path.append(comfyui_path)
+        print(f"'{comfyui_path}' added to sys.path")
+
+
+def add_extra_model_paths() -> None:
+    """
+    Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
+    """
+    try:
+        from main import load_extra_path_config
+    except ImportError:
+        print(
+            "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
+        )
+        from utils.extra_config import load_extra_path_config
+
+    extra_model_paths = find_path("extra_model_paths.yaml")
+
+    if extra_model_paths is not None:
+        load_extra_path_config(extra_model_paths)
+    else:
+        print("Could not find the extra_model_paths config file.")
+
+
+def import_custom_nodes() -> None:
+    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
+
+    This function sets up a new asyncio event loop, initializes the PromptServer,
+    creates a PromptQueue, and initializes the custom nodes.
+    """
+    import asyncio
+    import execution
+    from nodes import init_extra_nodes
+    import server
+
+    # Creating a new event loop and setting it as the default loop
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    # Creating an instance of PromptServer with the loop
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+
+    # Initializing custom nodes
+    init_extra_nodes()
+
+
+
+add_comfyui_directory_to_sys_path()
+import_custom_nodes()
+# add_extra_model_paths()
+dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+dualcliploader_11 = dualcliploader.load_clip(
+clip_name1="t5xxl_fp16.safetensors",
+clip_name2="clip_l.safetensors",
+type="flux",
+device="default",
+)
+
+cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+cliptextencode_6 = cliptextencode.encode(
+text="Photo on a small glass panel. Color. Photo of trees with a body of water in the front and moutain in the background.",
+clip=get_value_at_index(dualcliploader_11, 0),
+)
+
+vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+vaeloader_10 = vaeloader.load_vae(vae_name="ae.safetensors")
+
+unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+unetloader_12 = unetloader.load_unet(
+unet_name="flux1-dev.safetensors", weight_dtype="default"
+)
+
+ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
+ksamplerselect_16 = ksamplerselect.get_sampler(sampler_name="dpmpp_2m")
+
+randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
+randomnoise_25 = randomnoise.get_noise(noise_seed='42')
+
+loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
+loraloadermodelonly_72 = loraloadermodelonly.load_lora_model_only(
+lora_name='lora_weights/lora_weight_rank_32_alpha_32.safetensors',
+strength_model=1,
+model=get_value_at_index(unetloader_12, 0),
+)
+
+cr_sdxl_aspect_ratio = NODE_CLASS_MAPPINGS["CR SDXL Aspect Ratio"]()
+cr_sdxl_aspect_ratio_85 = cr_sdxl_aspect_ratio.Aspect_Ratio(
+width=1024,
+height=1024,
+aspect_ratio="4:3 landscape 1152x896",
+swap_dimensions="Off",
+upscale_factor=1.5,
+batch_size=1,
+)
+
+modelsamplingflux = NODE_CLASS_MAPPINGS["ModelSamplingFlux"]()
+fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
+basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]()
+basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
+samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
+vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+
+
+def model_loading():
+    # loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
+    # loraloadermodelonly_72 = loraloadermodelonly.load_lora_model_only(
+    #     lora_name=lora_weight_path,
+    #     strength_model=1,
+    #     model=get_value_at_index(unetloader_12, 0),
+    # )
+    model_loaders = [dualcliploader_11, vaeloader_10, unetloader_12, loraloadermodelonly_72]
+    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)
+]
+    #Load the models
+    # model_management.load_models_gpu(valid_models)
+    
+
+def generate_image(prompt,
+        height,
+        width,
+        guidance_scale,
+        aspect_ratio,
+        seed,
+        num_inference_steps,
+        ):
+    cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+    cliptextencode_6 = cliptextencode.encode(
+        text=prompt,
+        clip=get_value_at_index(dualcliploader_11, 0),
+    )
+    cr_sdxl_aspect_ratio = NODE_CLASS_MAPPINGS["CR SDXL Aspect Ratio"]()
+    cr_sdxl_aspect_ratio_85 = cr_sdxl_aspect_ratio.Aspect_Ratio(
+        width=width,
+        height=height,
+        aspect_ratio=aspect_ratio,
+        swap_dimensions="Off",
+        upscale_factor=1.5,
+        batch_size=1,
+    )
+    with torch.inference_mode():
+        for q in range(1):
+            modelsamplingflux_61 = modelsamplingflux.patch(
+                max_shift=1.15,
+                base_shift=0.5,
+                width=get_value_at_index(cr_sdxl_aspect_ratio_85, 0),
+                height=get_value_at_index(cr_sdxl_aspect_ratio_85, 1),
+                model=get_value_at_index(loraloadermodelonly_72, 0),
+            )
+
+            fluxguidance_60 = fluxguidance.append(
+                guidance=guidance_scale, conditioning=get_value_at_index(cliptextencode_6, 0)
+            )
+
+            basicguider_22 = basicguider.get_guider(
+                model=get_value_at_index(modelsamplingflux_61, 0),
+                conditioning=get_value_at_index(fluxguidance_60, 0),
+            )
+
+            basicscheduler_17 = basicscheduler.get_sigmas(
+                scheduler="sgm_uniform",
+                steps=num_inference_steps,
+                denoise=1,
+                model=get_value_at_index(modelsamplingflux_61, 0),
+            )
+
+            samplercustomadvanced_13 = samplercustomadvanced.sample(
+                noise=get_value_at_index(randomnoise_25, 0),
+                guider=get_value_at_index(basicguider_22, 0),
+                sampler=get_value_at_index(ksamplerselect_16, 0),
+                sigmas=get_value_at_index(basicscheduler_17, 0),
+                latent_image=get_value_at_index(cr_sdxl_aspect_ratio_85, 4),
+            )
+
+            vaedecode_8 = vaedecode.decode(
+                samples=get_value_at_index(samplercustomadvanced_13, 0),
+                vae=get_value_at_index(vaeloader_10, 0),
+            )
+
+            # saveimage_9 = saveimage.save_images(
+            #     filename_prefix="MarkuryFLUX", images=get_value_at_index(vaedecode_8, 0)
+            # )
+            return get_value_at_index(vaedecode_8, 0), seed