import gradio as gr import numpy as np import random import os import spaces #[uncomment to use ZeroGPU] from diffusers import DiffusionPipeline import torch from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler from huggingface_hub import hf_hub_download from safetensors.torch import load_file import sys sys.path.append('.') from utils.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV model_repo_id = "stabilityai/stable-diffusion-xl-base-1.0" repo_name = "tianweiy/DMD2" ckpt_name = "dmd2_sdxl_4step_unet_fp16.bin" device = "cuda" if torch.cuda.is_available() else "cpu" if torch.cuda.is_available(): torch_dtype = torch.bfloat16 else: torch_dtype = torch.float32 # Load model. unet = UNet2DConditionModel.from_config(model_repo_id, subfolder="unet").to(device, torch_dtype) unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name))) pipe = DiffusionPipeline.from_pretrained(model_repo_id, unet=unet, torch_dtype=torch_dtype).to(device) pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) unet = pipe.unet ## Change these parameters based on how you trained your sliderspace sliders train_method = 'xattn-strict' rank = 1 alpha =1 networks = {} modules = DEFAULT_TARGET_REPLACE modules += UNET_TARGET_REPLACE_MODULE_CONV for i in range(1): networks[i] = LoRANetwork( unet, rank=int(rank), multiplier=1.0, alpha=int(alpha), train_method=train_method, fast_init=True, ).to(device, dtype=torch_dtype) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 @spaces.GPU #[uncomment to use ZeroGPU] def infer( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, slider_space, discovered_directions, slider_scale, progress=gr.Progress(track_tqdm=True), ): if randomize_seed: seed = random.randint(0, MAX_SEED) sliderspace_path = f"sliderspace_weights/{slider_space}/slider_{int(discovered_directions[-1])-1}.pt" for net in networks: networks[net].load_state_dict(torch.load(sliderspace_path)) for net in networks: networks[net].set_lora_slider(slider_scale) with networks[0]: pass # original image generator = torch.Generator().manual_seed(seed) image = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, ).images[0] # edited image generator = torch.Generator().manual_seed(seed) with networks[0]: slider_image = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, ).images[0] return image, slider_image, seed examples = [ "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "An astronaut riding a green horse", "A delicious ceviche cheesecake slice", ] css = """ #col-container { margin: 0 auto; max-width: 640px; } """ ORIGINAL_SPACE_ID = 'baulab/SliderSpace' SPACE_ID = os.getenv('SPACE_ID') SHARED_UI_WARNING = f'''## You can duplicate and use it with a gpu with at least 24GB, or clone this repository to run on your own machine.
Duplicate Space
''' with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown(" # SliderSpace: Decomposing Visual Capabilities of Diffusion Models") # Adding links under the title gr.Markdown(""" 🔗 [Project Page](https://sliderspace.baulab.info) | 💻 [GitHub Code](https://github.com/rohitgandikota/sliderspace) """) with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0, variant="primary") # New dropdowns side by side with gr.Row(): slider_space = gr.Dropdown( choices= [ "alien", "ancient ruins", "animal", "bike", "car", "Citadel", "coral", "cowboy", "face", "futuristic cities", "monster", "mystical creature", "planet", "plant", "robot", "sculpture", "spaceship", "statue", "studio", "video game", "wizard" ], label="SliderSpace", value="spaceship" ) discovered_directions = gr.Dropdown( choices=[f"Slider {i}" for i in range(1, 11)], label="Discovered Directions", value="Slider 1" ) slider_scale = gr.Slider( label="Slider Scale", minimum=-10, maximum=10, step=0.1, value=1, ) with gr.Row(): result = gr.Image(label="Original Image", show_label=True) slider_result = gr.Image(label="Discovered Edit Direction", show_label=True) with gr.Accordion("Advanced Settings", open=False): negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=False, ) seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, # Replace with defaults that work for your model ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, # Replace with defaults that work for your model ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=2.0, step=0.1, value=0.0, # Replace with defaults that work for your model ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=4, # Replace with defaults that work for your model ) # gr.Examples(examples=examples, inputs=[prompt]) gr.on( triggers=[run_button.click, prompt.submit], fn=infer, inputs=[ prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, slider_space, discovered_directions, slider_scale ], outputs=[result, slider_result, seed], ) if __name__ == "__main__": demo.launch(share=True)