import gradio as gr import torch import os from utils import call from diffusers.pipelines import StableDiffusionXLPipeline StableDiffusionXLPipeline.__call__ = call import os from trainscripts.textsliders.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV from trainscripts.textsliders.demotrain import train_xl os.environ['CURL_CA_BUNDLE'] = '' model_map = {'Age' : 'models/age.pt', 'Chubby': 'models/chubby.pt', 'Muscular': 'models/muscular.pt', 'Wavy Eyebrows': 'models/eyebrows.pt', 'Small Eyes': 'models/eyesize.pt', 'Long Hair' : 'models/longhair.pt', 'Curly Hair' : 'models/curlyhair.pt', 'Smiling' : 'models/smiling.pt', 'Pixar Style' : 'models/pixar_style.pt', 'Sculpture Style': 'models/sculpture_style.pt', 'Repair Images': 'models/repair_slider.pt', 'Fix Hands': 'models/fix_hands.pt', } ORIGINAL_SPACE_ID = 'baulab/ConceptSliders' SPACE_ID = os.getenv('SPACE_ID') SHARED_UI_WARNING = f'''## Attention - Training does not work in this shared UI. You can either duplicate and use it with a gpu with at least 40GB, or clone this repository to run on your own machine.
Duplicate Space
''' class Demo: def __init__(self) -> None: self.training = False self.generating = False self.device = 'cuda' self.weight_dtype = torch.float16 self.pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=self.weight_dtype).to(self.device) with gr.Blocks() as demo: self.layout() demo.queue().launch(share=True, max_threads=3) def layout(self): with gr.Row(): if SPACE_ID == ORIGINAL_SPACE_ID: self.warning = gr.Markdown(SHARED_UI_WARNING) with gr.Row(): with gr.Tab("Test") as inference_column: with gr.Row(): self.explain_infr = gr.Markdown(value='This is a demo of [Concept Sliders: LoRA Adaptors for Precise Control in Diffusion Models](https://sliders.baulab.info/). To try out a model that can control a particular concept, select a model and enter any prompt, choose a seed, and finally choose the SDEdit timestep for structural preservation. Higher SDEdit timesteps results in more structural change. For example, if you select the model "Surprised Look" you can generate images for the prompt "A picture of a person, realistic, 8k" and compare the slider effect to the image generated by original model. We have also provided several other pre-fine-tuned models like "repair" sliders to repair flaws in SDXL generated images (Check out the "Pretrained Sliders" drop-down). You can also train and run your own custom sliders. Check out the "train" section for custom concept slider training.') with gr.Row(): with gr.Column(scale=1): self.prompt_input_infr = gr.Text( placeholder="Enter prompt...", label="Prompt", info="Prompt to generate" ) with gr.Row(): self.model_dropdown = gr.Dropdown( label="Pretrained Sliders", choices= list(model_map.keys()), value='Age', interactive=True ) self.seed_infr = gr.Number( label="Seed", value=12345 ) self.slider_scale_infr = gr.Slider( -6, 6, label="Slider Scale", value=2, info="Larger slider scale result in stronger edit" ) self.start_noise_infr = gr.Slider( 600, 900, value=750, label="SDEdit Timestep", info="Choose smaller values for more structural preservation" ) with gr.Column(scale=2): self.infr_button = gr.Button( value="Generate", interactive=True ) with gr.Row(): self.image_new = gr.Image( label="Slider", interactive=False, type='pil', ) self.image_orig = gr.Image( label="Original SD", interactive=False, type='pil', ) with gr.Tab("Train") as training_column: with gr.Row(): self.explain_train= gr.Markdown(value='In this part you can train a concept slider for Stable Diffusion XL. Enter a target concept you wish to make an edit on. Next, enter a enhance prompt of the attribute you wish to edit (for controlling age of a person, enter "person, old"). Then, type the supress prompt of the attribute (for our example, enter "person, young"). Then press "train" button. With default settings, it takes about 15 minutes to train a slider; then you can try inference above or download the weights. Code and details are at [github link](https://github.com/rohitgandikota/sliders).') with gr.Row(): with gr.Column(scale=3): self.target_concept = gr.Text( placeholder="Enter target concept to make edit on ...", label="Prompt of concept on which edit is made", info="Prompt corresponding to concept to edit" ) self.positive_prompt = gr.Text( placeholder="Enter the enhance prompt for the edit...", label="Prompt to enhance", info="Prompt corresponding to concept to enhance" ) self.negative_prompt = gr.Text( placeholder="Enter the suppress prompt for the edit...", label="Prompt to suppress", info="Prompt corresponding to concept to supress" ) self.rank = gr.Number( value=4, label="Rank of the Slider", info='Slider Rank to train' ) self.iterations_input = gr.Number( value=1000, precision=0, label="Iterations", info='iterations used to train' ) self.lr_input = gr.Number( value=2e-4, label="Learning Rate", info='Learning rate used to train' ) with gr.Column(scale=1): self.train_status = gr.Button(value='', variant='primary', interactive=False) self.train_button = gr.Button( value="Train", ) self.download = gr.Files() self.infr_button.click(self.inference, inputs = [ self.prompt_input_infr, self.seed_infr, self.start_noise_infr, self.slider_scale_infr, self.model_dropdown ], outputs=[ self.image_new, self.image_orig ] ) self.train_button.click(self.train, inputs = [ self.target_concept, self.positive_prompt, self.negative_prompt, self.rank, self.iterations_input, self.lr_input ], outputs=[self.train_button, self.train_status, self.download, self.model_dropdown] ) def train(self, target_concept,positive_prompt, negative_prompt, rank, iterations_input, lr_input, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)): randn = torch.randint(1, 10000000, (1,)).item() save_name = f'{randn}_{target_concept.replace(',','').replace(' ','').replace('.','')[:10]}_{positive_prompt.replace(',','').replace(' ','').replace('.','')[:10]}' save_name += f'_alpha-{1}' save_name += f'_noxattn' save_name += f'_rank_{rank}.pt' if self.training: return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()] self.training = True train_xl(target, postive, negative, lr, iterations, config_file, rank, device, attributes) self.training = False torch.cuda.empty_cache() model_map['Custom Slider'] = f'models/{save_name}' return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom slider in the "Test" tab'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom Slider')] # if train_method == 'ESD-x': # modules = ".*attn2$" # frozen = [] # elif train_method == 'ESD-u': # modules = "unet$" # frozen = [".*attn2$", "unet.time_embedding$", "unet.conv_out$"] # elif train_method == 'ESD-self': # modules = ".*attn1$" # frozen = [] # # save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}.pt" # self.training = True # train(prompt, modules, frozen, iterations, neg_guidance, lr, save_path) # self.training = False # torch.cuda.empty_cache() # model_map['Custom'] = save_path # return [None, None, None, None] def inference(self, prompt, seed, start_noise, scale, model_name, pbar = gr.Progress(track_tqdm=True)): seed = seed or 12345 generator = torch.manual_seed(seed) model_path = model_map[model_name] unet = self.pipe.unet network_type = "c3lier" if 'full' in model_path: train_method = 'full' elif 'noxattn' in model_path: train_method = 'noxattn' elif 'xattn' in model_path: train_method = 'xattn' network_type = 'lierla' else: train_method = 'noxattn' modules = DEFAULT_TARGET_REPLACE if network_type == "c3lier": modules += UNET_TARGET_REPLACE_MODULE_CONV name = os.path.basename(model_path) rank = 4 alpha = 1 if rank in model_path: rank = int(model_path.split('_')[-1].replace('.pt','')) # if 'rank4' in model_path: # rank = 4 # if 'rank8' in model_path: # rank = 8 if 'alpha1' in model_path: alpha = 1.0 network = LoRANetwork( unet, rank=rank, multiplier=1.0, alpha=alpha, train_method=train_method, ).to(self.device, dtype=self.weight_dtype) network.load_state_dict(torch.load(model_path)) generator = torch.manual_seed(seed) edited_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=50, generator=generator, network=network, start_noise=start_noise, scale=scale, unet=unet).images[0] generator = torch.manual_seed(seed) original_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=50, generator=generator, network=network, start_noise=start_noise, scale=0, unet=unet).images[0] del unet, network unet = None network = None pipe = None torch.cuda.empty_cache() return edited_image, original_image demo = Demo()