Style: Run Black formatter

app.py

@@ -12,7 +12,7 @@ from PIL import Image, ImageChops
 from huggingface_hub import hf_hub_download
 
 # Setup ComfyUI if not already set up
-#if not os.path.exists("ComfyUI"):
+# if not os.path.exists("ComfyUI"):
 #    print("Setting up ComfyUI...")
 #    subprocess.run(["bash", "setup_comfyui.sh"], check=True)
 
@@ -21,25 +21,63 @@ os.makedirs("output", exist_ok=True)
 
 # Download models if not already present
 print("Checking and downloading models...")
-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="black-forest-labs/FLUX.1-Canny-dev", filename="flux1-canny-dev.safetensors", local_dir="models/controlnet")
-hf_hub_download(repo_id="XLabs-AI/flux-controlnet-collections", filename="flux-canny-controlnet-v3.safetensors", local_dir="models/controlnet")
-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")
-t5_path = hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir="models/text_encoders/t5")
+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="black-forest-labs/FLUX.1-Canny-dev",
+    filename="flux1-canny-dev.safetensors",
+    local_dir="models/controlnet",
+)
+hf_hub_download(
+    repo_id="XLabs-AI/flux-controlnet-collections",
+    filename="flux-canny-controlnet-v3.safetensors",
+    local_dir="models/controlnet",
+)
+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",
+)
+t5_path = hf_hub_download(
+    repo_id="comfyanonymous/flux_text_encoders",
+    filename="t5xxl_fp16.safetensors",
+    local_dir="models/text_encoders/t5",
+)
 
 # Import required functions and setup ComfyUI path
 import folder_paths
 
+
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     try:
         return obj[index]
     except KeyError:
         return obj["result"][index]
 
+
 def find_path(name: str, path: str = None) -> str:
     if path is None:
         path = os.getcwd()
@@ -52,12 +90,14 @@ def find_path(name: str, path: str = None) -> str:
         return None
     return find_path(name, parent_directory)
 
+
 def add_comfyui_directory_to_sys_path() -> None:
     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:
     try:
         from main import load_extra_path_config
@@ -69,27 +109,30 @@ def add_extra_model_paths() -> None:
     else:
         print("Could not find the extra_model_paths config file.")
 
+
 # Initialize paths
 add_comfyui_directory_to_sys_path()
 add_extra_model_paths()
 
+
 def import_custom_nodes() -> None:
     import asyncio
     import execution
     from nodes import init_extra_nodes
     import server
-    
+
     # Create a new event loop if running in a new thread
     try:
         loop = asyncio.get_event_loop()
     except RuntimeError:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-    
+
     server_instance = server.PromptServer(loop)
     execution.PromptQueue(server_instance)
     init_extra_nodes()
 
+
 # Import all necessary nodes
 print("Importing ComfyUI nodes...")
 try:
@@ -108,7 +151,7 @@ try:
         UNETLoader,
         CLIPTextEncode,
     )
-    
+
     # Initialize all constant nodes and models in global context
     import_custom_nodes()
 except Exception as e:
@@ -156,7 +199,9 @@ SAMPLER = ksamplerselect.get_sampler(sampler_name="euler")
 
 # Initialize depth model
 cr_clip_input_switch = NODE_CLASS_MAPPINGS["CR Clip Input Switch"]()
-downloadandloaddepthanythingv2model = NODE_CLASS_MAPPINGS["DownloadAndLoadDepthAnythingV2Model"]()
+downloadandloaddepthanythingv2model = NODE_CLASS_MAPPINGS[
+    "DownloadAndLoadDepthAnythingV2Model"
+]()
 DEPTH_MODEL = downloadandloaddepthanythingv2model.loadmodel(
     model="depth_anything_v2_vitl_fp32.safetensors"
 )
@@ -177,7 +222,7 @@ clipvisionencode = CLIPVisionEncode()
 stylemodelapplyadvanced = NODE_CLASS_MAPPINGS["StyleModelApplyAdvanced"]()
 emptylatentimage = EmptyLatentImage()
 basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]()
-basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()        
+basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
 randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
 samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
 vaedecode = VAEDecode()
@@ -193,18 +238,31 @@ from comfy import model_management
 model_loaders = [CLIP_MODEL, VAE_MODEL, UNET_MODEL, CLIP_VISION_MODEL]
 
 print("Loading models to GPU...")
-model_management.load_models_gpu([
-    loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders
-])
+model_management.load_models_gpu(
+    [
+        loader[0].patcher if hasattr(loader[0], "patcher") else loader[0]
+        for loader in model_loaders
+    ]
+)
 
 print("Setup complete!")
 
+
 @spaces.GPU
-def generate_image(prompt, structure_image, style_image, depth_strength=15, canny_strength=30, style_strength=0.5, steps=28, progress=gr.Progress(track_tqdm=True)):
+def generate_image(
+    prompt,
+    structure_image,
+    style_image,
+    depth_strength=15,
+    canny_strength=30,
+    style_strength=0.5,
+    steps=28,
+    progress=gr.Progress(track_tqdm=True),
+):
     """Main generation function that processes inputs and returns the path to the generated image."""
     timestamp = random.randint(10000, 99999)
     output_filename = f"flux_zen_{timestamp}.png"
-    
+
     with torch.inference_mode():
         # Set up CLIP
         clip_switch = cr_clip_input_switch.switch(
@@ -212,7 +270,7 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             clip1=get_value_at_index(CLIP_MODEL, 0),
             clip2=get_value_at_index(CLIP_MODEL, 0),
         )
-        
+
         # Encode text
         text_encoded = cliptextencode.encode(
             text=prompt,
@@ -222,10 +280,10 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             text="",
             clip=get_value_at_index(clip_switch, 0),
         )
-        
+
         # Process structure image
         structure_img = loadimage.load_image(image=structure_image)
-        
+
         # Resize image
         resized_img = imageresize.execute(
             width=get_value_at_index(CONST_1024, 0),
@@ -236,59 +294,59 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             multiple_of=16,
             image=get_value_at_index(structure_img, 0),
         )
-        
+
         # Get image size
         size_info = getimagesizeandcount.getsize(
             image=get_value_at_index(resized_img, 0)
         )
-        
+
         # Encode VAE
         vae_encoded = vaeencode.encode(
             pixels=get_value_at_index(size_info, 0),
             vae=get_value_at_index(VAE_MODEL, 0),
         )
-        
+
         # Process canny
         canny_processed = canny_prossessor.detect_edge(
-            image=get_value_at_index(size_info, 0), 
-            low_threshold=0.4, 
-            high_threshold=0.8
-            )
-        
-        #Apply canny Advanced
+            image=get_value_at_index(size_info, 0),
+            low_threshold=0.4,
+            high_threshold=0.8,
+        )
+
+        # Apply canny Advanced
         canny_conditions = controlNetApplyAdvanced.apply_controlnet(
-            positive=get_value_at_index(text_encoded, 0), 
-            negative=get_value_at_index(empty_text, 0), 
-            control_net=get_value_at_index(CANNY_XLABS_MODEL, 0), 
-            image=get_value_at_index(canny_processed, 0), 
-            strength=canny_strength, 
-            start_percent=0.0, 
-            end_percent=0.5, 
-            vae=get_value_at_index(VAE_MODEL, 0)
-            )
-        
+            positive=get_value_at_index(text_encoded, 0),
+            negative=get_value_at_index(empty_text, 0),
+            control_net=get_value_at_index(CANNY_XLABS_MODEL, 0),
+            image=get_value_at_index(canny_processed, 0),
+            strength=canny_strength,
+            start_percent=0.0,
+            end_percent=0.5,
+            vae=get_value_at_index(VAE_MODEL, 0),
+        )
+
         # Process depth
         depth_processed = depthanything_v2.process(
             da_model=get_value_at_index(DEPTH_MODEL, 0),
             images=get_value_at_index(size_info, 0),
         )
-        
+
         # Apply Flux guidance
         flux_guided = fluxguidance.append(
             guidance=depth_strength,
             conditioning=get_value_at_index(canny_conditions, 0),
         )
-        
+
         # Process style image
         style_img = loadimage.load_image(image=style_image)
-        
+
         # Encode style with CLIP Vision
         style_encoded = clipvisionencode.encode(
             crop="center",
             clip_vision=get_value_at_index(CLIP_VISION_MODEL, 0),
             image=get_value_at_index(style_img, 0),
         )
-        
+
         # Set up conditioning
         conditioning = instructpixtopixconditioning.encode(
             positive=get_value_at_index(flux_guided, 0),
@@ -296,7 +354,7 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             vae=get_value_at_index(VAE_MODEL, 0),
             pixels=get_value_at_index(depth_processed, 0),
         )
-        
+
         # Apply style
         style_applied = stylemodelapplyadvanced.apply_stylemodel(
             strength=style_strength,
@@ -304,20 +362,20 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             style_model=get_value_at_index(STYLE_MODEL, 0),
             clip_vision_output=get_value_at_index(style_encoded, 0),
         )
-        
+
         # Set up empty latent
         empty_latent = emptylatentimage.generate(
             width=get_value_at_index(resized_img, 1),
             height=get_value_at_index(resized_img, 2),
             batch_size=1,
         )
-        
+
         # Set up guidance
         guided = basicguider.get_guider(
             model=get_value_at_index(UNET_MODEL, 0),
             conditioning=get_value_at_index(style_applied, 0),
         )
-        
+
         # Set up scheduler
         schedule = basicscheduler.get_sigmas(
             scheduler="simple",
@@ -325,10 +383,10 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             denoise=1,
             model=get_value_at_index(UNET_MODEL, 0),
         )
-        
+
         # Generate random noise
         noise = randomnoise.get_noise(noise_seed=random.randint(1, 2**64))
-        
+
         # Sample
         sampled = samplercustomadvanced.sample(
             noise=get_value_at_index(noise, 0),
@@ -337,25 +395,25 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
             sigmas=get_value_at_index(schedule, 0),
             latent_image=get_value_at_index(empty_latent, 0),
         )
-        
+
         # Decode VAE
         decoded = vaedecode.decode(
             samples=get_value_at_index(sampled, 0),
             vae=get_value_at_index(VAE_MODEL, 0),
         )
-        
+
         # Create text node for prefix
         prefix = cr_text.text_multiline(text=f"flux_zen_{timestamp}")
-        
+
         # Use SaveImage node to save the image
         saved_data = saveimage.save_images(
             filename_prefix=get_value_at_index(prefix, 0),
             images=get_value_at_index(decoded, 0),
         )
-        
+
         try:
             saved_path = f"output/{saved_data['ui']['images'][0]['filename']}"
-                
+
             return saved_path
         except Exception as e:
             print(f"Error getting saved image path: {e}")
@@ -365,59 +423,56 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, cann
 
 with gr.Blocks(css="footer {visibility: hidden}") as app:
     gr.Markdown("# 🎨 FLUX Zen Style Depth+Canny")
-    gr.Markdown("Flux[dev] Redux + Flux[dev] Depth and XLabs Canny based on the space FLUX Style Shaping")
-    
+    gr.Markdown(
+        "Flux[dev] Redux + Flux[dev] Depth and XLabs Canny based on the space FLUX Style Shaping"
+    )
+
     with gr.Row():
         with gr.Column(scale=1):
             prompt_input = gr.Textbox(
-                label="Prompt", 
+                label="Prompt",
                 placeholder="Enter your prompt here...",
-                info="Describe the image you want to generate"
+                info="Describe the image you want to generate",
             )
             with gr.Row():
                 with gr.Column(scale=1):
                     structure_image = gr.Image(
-                        image_mode='RGB', 
-                        label="Structure Image", 
-                        type="filepath"
+                        image_mode="RGB", label="Structure Image", type="filepath"
                     )
                     depth_strength = gr.Slider(
-                        minimum=0, 
-                        maximum=50, 
-                        value=15, 
+                        minimum=0,
+                        maximum=50,
+                        value=15,
                         label="Depth Strength",
-                        info="Controls how much the depth map influences the result"
+                        info="Controls how much the depth map influences the result",
                     )
                     canny_strength = gr.Slider(
-                        minimum=0, 
-                        maximum=1.0, 
-                        value=0.30, 
+                        minimum=0,
+                        maximum=1.0,
+                        value=0.30,
                         label="Canny Strength",
-                        info="Controls how much the edge detection influences the result"
+                        info="Controls how much the edge detection influences the result",
                     )
                     steps = gr.Slider(
-                        minimum=10, 
-                        maximum=50, 
-                        value=28, 
+                        minimum=10,
+                        maximum=50,
+                        value=28,
                         label="Steps",
-                        info="More steps = better quality but slower generation"
+                        info="More steps = better quality but slower generation",
                     )
                 with gr.Column(scale=1):
-                    style_image = gr.Image(
-                        label="Style Image", 
-                        type="filepath"
-                    )
+                    style_image = gr.Image(label="Style Image", type="filepath")
                     style_strength = gr.Slider(
-                        minimum=0, 
-                        maximum=1, 
-                        value=0.5, 
+                        minimum=0,
+                        maximum=1,
+                        value=0.5,
                         label="Style Strength",
-                        info="Controls how much the style image influences the result"
+                        info="Controls how much the style image influences the result",
                     )
-            
+
             with gr.Row():
-                generate_btn = gr.Button("Generate",value=True, variant="primary")
-                
+                generate_btn = gr.Button("Generate", value=True, variant="primary")
+
         with gr.Column(scale=1):
             output_image = gr.Image(label="Generated Image")
 
@@ -441,11 +496,20 @@ with gr.Blocks(css="footer {visibility: hidden}") as app:
 
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt_input, structure_image, style_image, depth_strength, canny_strength, style_strength, steps],
-        outputs=[output_image]
+        inputs=[
+            prompt_input,
+            structure_image,
+            style_image,
+            depth_strength,
+            canny_strength,
+            style_strength,
+            steps,
+        ],
+        outputs=[output_image],
     )
-    
-    gr.Markdown("""
+
+    gr.Markdown(
+        """
     ## How to use
     1. Enter a prompt describing the image you want to generate
     2. Upload a structure image to provide the basic shape/composition
@@ -456,11 +520,12 @@ with gr.Blocks(css="footer {visibility: hidden}") as app:
     ## About
     This demo uses FLUX.1-Redux-dev for style transfer, FLUX.1-Depth-dev for depth-guided generation, 
     and XLabs Canny for edge detection and structure preservation.
-    """)
+    """
+    )
 
 if __name__ == "__main__":
     # Create an examples directory if it doesn't exist , for now it is empty
     os.makedirs("examples", exist_ok=True)
-    
+
     # Launch the app
-    app.launch(share=True)
+    app.launch(share=True)