ControlNet inference

apt-get update && apt-get install ffmpeg libsm6 libxext6 -y

python src/maxdiffusion/controlnet/generate_controlnet_replicated.py

config.prompt = "a photograph of a modern building"
config.negative_prompt = "blurry, distorted"
config.controlnet_image = "https://example.com/edge_map.png"
config.controlnet_model_name_or_path = "lllyasviel/sd-controlnet-canny"
config.controlnet_from_pt = True
config.controlnet_conditioning_scale = 1.0
config.num_inference_steps = 50
config.per_device_batch_size = 1
config.seed = 0

# Load ControlNet
controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
    config.controlnet_model_name_or_path, 
    from_pt=config.controlnet_from_pt, 
    dtype=jnp.float32
)

# Load pipeline with ControlNet
pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
    config.pretrained_model_name_or_path, 
    controlnet=controlnet, 
    revision=config.revision, 
    dtype=jnp.float32
)
params["controlnet"] = controlnet_params

# Prepare inputs
prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
processed_image = pipe.prepare_image_inputs([canny_image] * num_samples)

# Replicate and shard
p_params = replicate(params)
prompt_ids = shard(prompt_ids)
negative_prompt_ids = shard(negative_prompt_ids)
processed_image = shard(processed_image)

# Generate
output = pipe(
    prompt_ids=prompt_ids,
    image=processed_image,
    params=p_params,
    prng_seed=rng,
    num_inference_steps=config.num_inference_steps,
    neg_prompt_ids=negative_prompt_ids,
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    jit=True,
).images

python src/maxdiffusion/controlnet/generate_controlnet_sdxl_replicated.py

image = load_image(config.controlnet_image)
image = np.array(image)

# Apply Canny edge detection
image = cv2.Canny(image, 100, 200)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
image = Image.fromarray(image)

controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
    config.controlnet_model_name_or_path, 
    from_pt=config.controlnet_from_pt, 
    dtype=config.activations_dtype
)

pipe, params = FlaxStableDiffusionXLControlNetPipeline.from_pretrained(
    config.pretrained_model_name_or_path, 
    controlnet=controlnet, 
    revision=config.revision, 
    dtype=config.activations_dtype
)

# Cast params to bfloat16 (except scheduler)
scheduler_state = params.pop("scheduler")
params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), params)
params["scheduler"] = scheduler_state
params["controlnet"] = controlnet_params

# Light control for more variation
config.controlnet_conditioning_scale = 0.5

# Strong control for precise structure
config.controlnet_conditioning_scale = 1.5

config.controlnet_image = "/path/to/custom_edge_map.png"

num_samples = jax.device_count() * config.per_device_batch_size
rng = jax.random.split(rng, jax.device_count())

# Replicate params across devices
p_params = replicate(params)

# Shard inputs
prompt_ids = shard(prompt_ids)
processed_image = shard(processed_image)

for i, image in enumerate(output_images):
  image.save(f"generated_image_{i}.png")

config.prompt = "a modern glass office building, blue sky, photorealistic"
config.negative_prompt = "blurry, cartoon, painting"
config.controlnet_image = "building_edges.png"
config.controlnet_conditioning_scale = 1.2

config.prompt = "portrait of a woman in a red dress, studio lighting"
config.negative_prompt = "deformed, blurry"
config.controlnet_model_name_or_path = "lllyasviel/sd-controlnet-openpose"
config.controlnet_image = "pose_skeleton.png"
config.controlnet_conditioning_scale = 1.0

config.prompt = "mountain landscape at sunset, dramatic clouds"
config.negative_prompt = "flat, boring"
config.controlnet_model_name_or_path = "lllyasviel/sd-controlnet-depth"
config.controlnet_image = "depth_map.png"
config.controlnet_conditioning_scale = 0.8