Update README.md

README.md

@@ -6,40 +6,42 @@ This is a custom block designed to extract depth maps from input images using th
 
 ```python
 import torch
-from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks
-from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS, CONTROLNET_BLOCKS
+from diffusers import ModularPipeline, ComponentsManager, ModularPipelineBlocks
 from diffusers.utils import load_image
 
-# fetch the depth processor block that will create our depth map
-depth_processor_block = ModularPipelineBlocks.from_pretrained("diffusers/depth-processor-custom-block", trust_remote_code=true)
+# Use ComponentsManager to enable auto CPU offloading for memory efficiency
+manager = ComponentsManager()
+manager.enable_auto_cpu_offload(device="cuda:0")
 
-my_blocks = TEXT2IMAGE_BLOCKS.copy()
-my_blocks.insert("depth_processor", depth_processor_block, 1)
 
-# replace text to image denoise block with controlnet denoise block
-my_blocks.sub_blocks["denoise"] = CONTROLNET_BLOCKS["denoise"]
+# Initialize pipeline
+pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image", components_manager=manager)
 
-# create our initial set of controlnet blocks
-blocks = SequentialPipelineBlocks.from_blocks_dict(my_blocks)
-
-repo_id = "diffusers/modular-stable-diffusion-xl-base-1.0"
-
-# Initialize the pipeline object we can use to run our blocks
-pipe = blocks.init_pipeline(repo_id)
+# Insert a depth processing block
+blocks = pipe.blocks.get_workflow("controlnet_text2image")
+depth_block = ModularPipelineBlocks.from_pretrained(
+    "diffusers/depth-processor-custom-block",
+    trust_remote_code=True,
+)
+blocks.sub_blocks.insert("depth", depth_block, 0)
 
-# Load model component weights
-pipe.load_components(torch_dtype=torch.float16, device_map="cuda")
+# Reinitialize the pipeline for ControlNet
+pipe = blocks.init_pipeline("Qwen/Qwen-Image", components_manager=manager)
+pipe.load_components(torch_dtype=torch.bfloat16)
 
-image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true")
-image = image.resize((1024, 1024))
+# Load the ControlNet model
+controlnet_spec = pipeline.get_component_spec("controlnet")
+controlnet_spec.pretrained_model_name_or_path = "InstantX/Qwen-Image-ControlNet-Union"
+controlnet = controlnet_spec.load(torch_dtype=torch.bfloat16)
+pipeline.update_components(controlnet=controlnet)
 
-prompt = ["A red car"]
+# Infer
+prompt = "cat wizard with red hat, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney"
+image = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/cute_cat.png?raw=true")
 
-output = pipe(
+output = pipeline(
     prompt=prompt,
     image=image,
-    num_inference_steps=35,
-    guidance_scale=7.5,
-    output_type="pil",
-)
+).images[0]
+output
 ```