multimodalart/ctrl-x
8
1from argparse import ArgumentParser2 3from diffusers import DDIMScheduler, StableDiffusionXLImg2ImgPipeline4import gradio as gr5import torch6import yaml7 8from ctrl_x.pipelines.pipeline_sdxl import CtrlXStableDiffusionXLPipeline9from ctrl_x.utils import *10from ctrl_x.utils.sdxl import *11 12import spaces13 14 15parser = ArgumentParser()16parser.add_argument("-m", "--model", type=str, default=None) # Optionally, load model checkpoint from single file17args = parser.parse_args()18 19torch.backends.cudnn.enabled = False # Sometimes necessary to suppress CUDNN_STATUS_NOT_SUPPORTED20 21torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float3222 23model_id_or_path = "stabilityai/stable-diffusion-xl-base-1.0"24refiner_id_or_path = "stabilityai/stable-diffusion-xl-refiner-1.0"25device = "cuda" if torch.cuda.is_available() else "cpu"26#variant = "fp16" if device == "cuda" else "fp32"27 28scheduler = DDIMScheduler.from_config(model_id_or_path, subfolder="scheduler") # TODO: Support other schedulers29if args.model is None:30 pipe = CtrlXStableDiffusionXLPipeline.from_pretrained(31 model_id_or_path, scheduler=scheduler, torch_dtype=torch_dtype, use_safetensors=True32 )33else:34 print(f"Using weights {args.model} for SDXL base model.")35 pipe = CtrlXStableDiffusionXLPipeline.from_single_file(args.model, scheduler=scheduler, torch_dtype=torch_dtype)36refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(37 refiner_id_or_path, scheduler=scheduler, text_encoder_2=pipe.text_encoder_2, vae=pipe.vae,38 torch_dtype=torch_dtype, use_safetensors=True,39)40 41if torch.cuda.is_available():42 pipe = pipe.to("cuda")43 refiner = refiner.to("cuda")44 45 46def get_control_config(structure_schedule, appearance_schedule):47 s = structure_schedule48 a = appearance_schedule49 50 control_config =\51f"""control_schedule:52 # structure_conv structure_attn appearance_attn conv/attn53 encoder: # (num layers)54 0: [[ ], [ ], [ ]] # 2/055 1: [[ ], [ ], [{a}, {a} ]] # 2/256 2: [[ ], [ ], [{a}, {a} ]] # 2/257 middle: [[ ], [ ], [ ]] # 2/158 decoder:59 0: [[{s} ], [{s}, {s}, {s}], [0.0, {a}, {a}]] # 3/360 1: [[ ], [ ], [{a}, {a} ]] # 3/361 2: [[ ], [ ], [ ]] # 3/062 63control_target:64 - [output_tensor] # structure_conv choices: {{hidden_states, output_tensor}}65 - [query, key] # structure_attn choices: {{query, key, value}}66 - [before] # appearance_attn choices: {{before, value, after}}67 68self_recurrence_schedule:69 - [0.1, 0.5, 2] # format: [start, end, num_recurrence]"""70 71 return control_config72 73 74css = """75.config textarea {font-family: monospace; font-size: 80%; white-space: pre}76.mono {font-family: monospace}77"""78 79title = """80<div style="display: flex; align-items: center; justify-content: center;margin-bottom: -15px">81 <h1 style="margin-left: 12px;text-align: center;display: inline-block">82 Ctrl-X: Controlling Structure and Appearance for Text-To-Image Generation Without Guidance83 </h1>84 <h3 style="display: inline-block; margin-left: 10px; margin-top: 7.5px; font-weight: 500">85 SDXL v1.086 </h3>87</div>88<div style="display: flex; align-items: center; justify-content: center;margin-bottom: 25px">89 <h3 style="text-align: center">90 [<a href="https://genforce.github.io/ctrl-x/">Page</a>]91 92 [<a href="https://arxiv.org/abs/2406.07540">Paper</a>]93 94 [<a href="https://github.com/genforce/ctrl-x">Code</a>]95 </h3>96</div>97"""98description = """<div>99 <p>100 <b>Ctrl-X</b> is a simple training-free and guidance-free framework for text-to-image (T2I) generation with 101 structure and appearance control. Given structure and appearance images, Ctrl-X designs feedforward structure 102 control to enable structure alignment with the arbitrary structure image and semantic-aware appearance transfer 103 to facilitate the appearance transfer from the appearance image.104 </p>105 <p>106 Here are some notes and tips for this demo:107 </p>108 <ul>109 <li> On input images:110 <ul>111 <li>112 If both the structure and appearance images are provided, then Ctrl-X does <i>structure and 113 appearance</i> control.114 </li>115 <li>116 If only the structure image is provided, then Ctrl-X does <i>structure-only</i> control and the 117 appearance image is jointly generated with the output image.118 </li>119 <li>120 Similarly, if only the appearance image is provided, then Ctrl-X does <i>appearance-only</i> 121 control.122 </li>123 </ul>124 </li>125 <li> On prompts:126 <ul>127 <li>128 Though the output prompt can affect the output image to a noticeable extent, the "accuracy" of the 129 structure and appearance prompts are not impactful to the final image.130 </li>131 <li>132 If the structure or appearance prompt is left blank, then it uses the (non-optional) output prompt 133 by default.134 </li>135 </ul>136 </li>137 <li> On control schedules:138 <ul>139 <li>140 When "Use advanced config" is <b>OFF</b>, the demo uses the structure guidance 141 (<span class="mono">structure_conv</span> and <span class="mono">structure_attn</span> 142 in the advanced config) and appearance guidance (<span class="mono">appearance_attn</span> in the 143 advanced config) sliders to change the control schedules.144 </li>145 <li>146 Otherwise, the demo uses "Advanced control config," which allows per-layer structure and 147 appearance schedule control, along with self-recurrence control. <i>This should be used 148 carefully</i>, and we recommend switching "Use advanced config" <b>OFF</b> in most cases. (For the 149 examples provided at the bottom of the demo, the advanced config uses the default schedules that 150 may not be the best settings for these examples.)151 </li>152 </ul>153 </li>154 </ul>155 <p>156 Have fun! :D157 </p>158</div>159"""160 161@spaces.GPU162def inference(163 structure_image,164 appearance_image,165 prompt,166 structure_prompt,167 appearance_prompt,168 positive_prompt="high quality",169 negative_prompt="ugly, blurry, dark, low res, unrealistic",170 guidance_scale=5.0,171 structure_guidance_scale=5.0,172 appearance_guidance_scale=5.0,173 num_inference_steps=28,174 eta=1.0,175 seed=42,176 width=1024,177 height=1024,178 structure_schedule=0.6,179 appearance_schedule=0.6,180 use_advanced_config=False,181 control_config="",182 progress=gr.Progress(track_tqdm=True)183):184 torch.manual_seed(seed)185 186 pipe.scheduler.set_timesteps(num_inference_steps, device=device)187 timesteps = pipe.scheduler.timesteps188 189 print(f"\nUsing the following control config (use_advanced_config={use_advanced_config}):")190 if not use_advanced_config:191 control_config = get_control_config(structure_schedule, appearance_schedule)192 print(control_config, end="\n\n")193 194 config = yaml.safe_load(control_config)195 register_control(196 model = pipe,197 timesteps = timesteps,198 control_schedule = config["control_schedule"],199 control_target = config["control_target"],200 )201 202 pipe.safety_checker = None203 pipe.requires_safety_checker = False204 205 self_recurrence_schedule = get_self_recurrence_schedule(config["self_recurrence_schedule"], num_inference_steps)206 207 pipe.set_progress_bar_config(desc="Ctrl-X inference")208 refiner.set_progress_bar_config(desc="Refiner")209 210 result, structure, appearance = pipe(211 prompt = prompt,212 structure_prompt = structure_prompt,213 appearance_prompt = appearance_prompt,214 structure_image = structure_image,215 appearance_image = appearance_image,216 num_inference_steps = num_inference_steps,217 negative_prompt = negative_prompt,218 positive_prompt = positive_prompt,219 height = height,220 width = width,221 guidance_scale = guidance_scale,222 structure_guidance_scale = structure_guidance_scale,223 appearance_guidance_scale = appearance_guidance_scale,224 eta = eta,225 output_type = "pil",226 return_dict = False,227 control_schedule = config["control_schedule"],228 self_recurrence_schedule = self_recurrence_schedule,229 )230 231 result_refiner = refiner(232 image = pipe.refiner_args["latents"],233 prompt = pipe.refiner_args["prompt"],234 negative_prompt = pipe.refiner_args["negative_prompt"],235 height = height,236 width = width,237 num_inference_steps = num_inference_steps,238 guidance_scale = guidance_scale,239 guidance_rescale = 0.7,240 num_images_per_prompt = 1,241 eta = eta,242 output_type = "pil",243 ).images244 del pipe.refiner_args245 246 return [result[0], result_refiner[0], structure[0], appearance[0]]247 248 249with gr.Blocks(theme=gr.themes.Default(), css=css, title="Ctrl-X (SDXL v1.0)") as app:250 gr.HTML(title)251 with gr.Accordion("Instructions", open=False):252 gr.HTML(description)253 with gr.Row():254 with gr.Column(scale=45):255 with gr.Group():256 kwargs = {} # {"width": 400, "height": 400}257 with gr.Row():258 structure_image = gr.Image(label="Upload structure image (optional)", type="pil", **kwargs)259 appearance_image = gr.Image(label="Upload appearance image (optional)", type="pil", **kwargs)260 with gr.Row():261 structure_prompt = gr.Textbox(label="Structure prompt (optional)", placeholder="Describes the structure image")262 appearance_prompt = gr.Textbox(label="Appearance prompt (optional)", placeholder="Describes the style image")263 with gr.Row():264 prompt = gr.Textbox(label="Output prompt", placeholder="Prompt which describes the output image")265 with gr.Row():266 positive_prompt = gr.Textbox(label="Positive prompt", value="high quality", placeholder="")267 negative_prompt = gr.Textbox(label="Negative prompt", value="ugly, blurry, dark, low res, unrealistic", placeholder="")268 with gr.Accordion("Advanced Options", open=False):269 with gr.Row():270 guidance_scale = gr.Slider(label="Target guidance scale", value=5.0, minimum=1, maximum=10)271 structure_guidance_scale = gr.Slider(label="Structure guidance scale", value=5.0, minimum=1, maximum=10)272 appearance_guidance_scale = gr.Slider(label="Appearance guidance scale", value=5.0, minimum=1, maximum=10)273 with gr.Row():274 num_inference_steps = gr.Slider(label="# inference steps", value=28, minimum=1, maximum=200, step=1)275 eta = gr.Slider(label="Eta (noise)", value=1.0, minimum=0, maximum=1.0, step=0.01)276 seed = gr.Slider(0, 2147483647, label="Seed", value=90095, step=1)277 with gr.Row():278 width = gr.Slider(label="Width", value=1024, minimum=256, maximum=2048, step=pipe.vae_scale_factor)279 height = gr.Slider(label="Height", value=1024, minimum=256, maximum=2048, step=pipe.vae_scale_factor)280 with gr.Row():281 structure_schedule = gr.Slider(label="Structure schedule", value=0.6, minimum=0.0, maximum=1.0, step=0.01, scale=2)282 appearance_schedule = gr.Slider(label="Appearance schedule", value=0.6, minimum=0.0, maximum=1.0, step=0.01, scale=2)283 use_advanced_config = gr.Checkbox(label="Use advanced config", value=False, scale=1)284 with gr.Row():285 control_config = gr.Textbox(286 label="Advanced control config", lines=20, value=get_control_config(0.6, 0.6), elem_classes=["config"], visible=False,287 )288 use_advanced_config.change(289 fn=lambda value: gr.update(visible=value), inputs=use_advanced_config, outputs=control_config,290 )291 with gr.Row():292 generate = gr.Button(value="Run")293 294 with gr.Column(scale=55):295 with gr.Group():296 with gr.Row():297 result_refiner = gr.Image(label="Output image w/ refiner", format="jpg", **kwargs)298 with gr.Row():299 result = gr.Image(label="Output image", format="jpg", **kwargs)300 structure_recon = gr.Image(label="Structure image", format="jpg", **kwargs)301 appearance_recon = gr.Image(label="Style image", format="jpg", **kwargs)302 303 inputs = [304 structure_image, appearance_image,305 prompt, structure_prompt, appearance_prompt,306 positive_prompt, negative_prompt,307 guidance_scale, structure_guidance_scale, appearance_guidance_scale,308 num_inference_steps, eta, seed,309 width, height,310 structure_schedule, appearance_schedule, use_advanced_config,311 control_config,312 ]313 outputs = [result, result_refiner, structure_recon, appearance_recon]314 315 generate.click(inference, inputs=inputs, outputs=outputs)316 317 examples = gr.Examples(318 [319 [320 "assets/images/horse__point_cloud.jpg",321 "assets/images/horse.jpg",322 "a photo of a horse standing on grass",323 "a 3D point cloud of a horse",324 "",325 ],326 [327 "assets/images/cat__mesh.jpg",328 "assets/images/tiger.jpg",329 "a photo of a tiger standing on snow",330 "a 3D mesh of a cat",331 "",332 ],333 [334 "assets/images/dog__sketch.jpg",335 "assets/images/squirrel.jpg",336 "a photo of a squirrel",337 "a sketch of a dog",338 "",339 ],340 [341 "assets/images/living_room__seg.jpg",342 "assets/images/van_gogh.jpg",343 "a Van Gogh painting of a living room",344 "a segmentation map of a living room",345 "",346 ],347 [348 "assets/images/bedroom__sketch.jpg",349 "assets/images/living_room_modern.jpg",350 "a sketch of a bedroom",351 "a photo of a modern bedroom during sunset",352 "",353 ],354 [355 "assets/images/running__pose.jpg",356 "assets/images/man_park.jpg",357 "a photo of a man running in a park",358 "a pose image of a person running",359 "",360 ],361 [362 "assets/images/fruit_bowl.jpg",363 "assets/images/grapes.jpg",364 "a photo of a bowl of grapes in the trees",365 "a photo of a bowl of fruits",366 "",367 ],368 [369 "assets/images/bear_avocado__spatext.jpg",370 None,371 "a realistic photo of a bear and an avocado in a forest",372 "a segmentation map of a bear and an avocado",373 "",374 ],375 [376 "assets/images/cat__point_cloud.jpg",377 None,378 "an embroidery of a white cat sitting on a rock under the night sky",379 "a 3D point cloud of a cat",380 "",381 ],382 [383 "assets/images/library__mesh.jpg",384 None,385 "a Polaroid photo of an old library, sunlight streaming in",386 "a 3D mesh of a library",387 "",388 ],389 [390 "assets/images/knight__humanoid.jpg",391 None,392 "a photo of a medieval soldier standing on a barren field, raining",393 "a 3D model of a person holding a sword and shield",394 "",395 ],396 [397 "assets/images/person__mesh.jpg",398 None,399 "a photo of a Karate man performing in a cyberpunk city at night",400 "a 3D mesh of a person",401 "",402 ],403 ],404 [405 structure_image,406 appearance_image,407 prompt,408 structure_prompt,409 appearance_prompt,410 ],411 examples_per_page=50,412 cache_examples="lazy",413 fn=inference,414 outputs=[result, result_refiner, structure_recon, appearance_recon]415 )416 417app.launch(debug=False, share=False)418