Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| import os | |
| from utils import call | |
| from diffusers import ( | |
| DDPMScheduler, | |
| DDIMScheduler, | |
| PNDMScheduler, | |
| LMSDiscreteScheduler, | |
| EulerAncestralDiscreteScheduler, | |
| EulerDiscreteScheduler, | |
| DPMSolverMultistepScheduler, | |
| ) | |
| 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', | |
| 'Surprised Look': 'models/suprised_look.pt', | |
| 'Smiling' : 'models/smiling.pt', | |
| 'Professional': 'models/professional.pt', | |
| 'Long Hair' : 'models/long_hair.pt', | |
| 'Curly Hair' : 'models/curlyhair.pt', | |
| 'Pixar Style' : 'models/pixar_style.pt', | |
| 'Sculpture Style': 'models/sculpture_style.pt', | |
| 'Clay Style': 'models/clay_style.pt', | |
| 'Repair Images': 'models/repair_slider.pt', | |
| 'Fix Hands': 'models/fix_hands.pt', | |
| 'Cluttered Room': 'models/cluttered_room.pt', | |
| 'Dark Weather': 'models/dark_weather.pt', | |
| 'Festive': 'models/festive.pt', | |
| 'Tropical Weather': 'models/tropical_weather.pt', | |
| 'Winter Weather': 'models/winter_weather.pt', | |
| 'Wavy Eyebrows': 'models/eyebrow.pt', | |
| 'Small Eyes (use scales -3, -1, 1, 3)': 'models/eyesize.pt', | |
| } | |
| ORIGINAL_SPACE_ID = 'baulab/ConceptSliders' | |
| SPACE_ID = os.getenv('SPACE_ID') | |
| SHARED_UI_WARNING = f'''## Attention - Training could be slow in this shared UI. You can alternatively duplicate and use it with a gpu with at least 40GB, or clone this repository to run on your own machine. | |
| <center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center> | |
| ''' | |
| class Demo: | |
| def __init__(self) -> None: | |
| self.training = False | |
| self.generating = False | |
| self.device = 'cuda' | |
| self.weight_dtype = torch.bfloat16 | |
| model_id = 'stabilityai/stable-diffusion-xl-base-1.0' | |
| pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=self.weight_dtype).to(self.device) | |
| pipe = None | |
| del pipe | |
| torch.cuda.empty_cache() | |
| model_id = "stabilityai/sdxl-turbo" | |
| self.current_model = 'SDXL Turbo' | |
| euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") | |
| self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=euler_anc, torch_dtype=self.weight_dtype).to(self.device) | |
| self.pipe.enable_xformers_memory_efficient_attention() | |
| self.guidance_scale = 1 | |
| self.num_inference_steps = 3 | |
| with gr.Blocks() as demo: | |
| self.layout() | |
| demo.queue(max_size=5).launch(share=True, max_threads=2) | |
| 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. <b>Current Inference is running on SDXL Turbo!</b>') | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| self.prompt_input_infr = gr.Text( | |
| placeholder="photo of a person, with bokeh street background, realistic, 8k", | |
| label="Prompt", | |
| info="Prompt to generate", | |
| value="photo of a person, with bokeh street background, realistic, 8k" | |
| ) | |
| 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=42753 | |
| ) | |
| self.slider_scale_infr = gr.Slider( | |
| -4, | |
| 4, | |
| label="Slider Scale", | |
| value=3, | |
| 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" | |
| ) | |
| self.model_type = gr.Dropdown( | |
| label="Model", | |
| choices= ['SDXL Turbo', 'SDXL'], | |
| value='SDXL Turbo', | |
| interactive=True | |
| ) | |
| with gr.Column(scale=2): | |
| self.infr_button = gr.Button( | |
| value="Generate", | |
| interactive=True | |
| ) | |
| with gr.Row(): | |
| self.image_orig = gr.Image( | |
| label="Original SD", | |
| interactive=False, | |
| type='pil', | |
| ) | |
| self.image_new = gr.Image( | |
| label=f"Concept Slider", | |
| 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 textual concept sliders for Stable Diffusion XL. Enter a target concept you wish to make an edit on (eg. person). 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 25 minutes to train a slider; then you can try inference above or download the weights. For faster training, please duplicate the repo and train with A100 or larger GPU. 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 (eg: 'person')", | |
| value = '' | |
| ) | |
| self.positive_prompt = gr.Text( | |
| placeholder="Enter the enhance prompt for the edit ...", | |
| label="Prompt to enhance", | |
| info="Prompt corresponding to concept to enhance (eg: 'person, old')", | |
| value = '' | |
| ) | |
| self.negative_prompt = gr.Text( | |
| placeholder="Enter the suppress prompt for the edit ...", | |
| label="Prompt to suppress", | |
| info="Prompt corresponding to concept to supress (eg: 'person, young')", | |
| value = '' | |
| ) | |
| self.attributes_input = gr.Text( | |
| placeholder="Enter the concepts to preserve (comma seperated). Leave empty if not required ...", | |
| label="Concepts to Preserve", | |
| info="Comma seperated concepts to preserve/disentangle (eg: 'male, female')", | |
| value = '' | |
| ) | |
| self.is_person = gr.Checkbox( | |
| label="Person", | |
| info="Are you training a slider for person?") | |
| self.rank = gr.Number( | |
| value=4, | |
| label="Rank of the Slider", | |
| info='Slider Rank to train' | |
| ) | |
| choices = ['xattn', 'noxattn'] | |
| self.train_method_input = gr.Dropdown( | |
| choices=choices, | |
| value='xattn', | |
| label='Train Method', | |
| info='Method of training. If [* xattn *] - loras will be on cross attns only. [* noxattn *] (official implementation) - all layers except cross attn', | |
| interactive=True | |
| ) | |
| self.iterations_input = gr.Number( | |
| value=500, | |
| precision=0, | |
| label="Iterations", | |
| info='iterations used to train - maximum of 1000' | |
| ) | |
| 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, | |
| self.model_type | |
| ], | |
| 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, | |
| self.attributes_input, | |
| self.is_person, | |
| self.train_method_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, attributes_input, is_person, train_method_input, pbar = gr.Progress(track_tqdm=True)): | |
| iterations_input = min(int(iterations_input),1000) | |
| if attributes_input == '': | |
| attributes_input = None | |
| print(target_concept, positive_prompt, negative_prompt, attributes_input, is_person) | |
| randn = torch.randint(1, 10000000, (1,)).item() | |
| save_name = f"{randn}_{positive_prompt.replace(',','').replace(' ','').replace('.','')[:20]}" | |
| save_name += f'_alpha-{1}' | |
| save_name += f'_{train_method_input}' | |
| save_name += f'_rank_{int(rank)}.pt' | |
| # if torch.cuda.get_device_properties(0).total_memory * 1e-9 < 40: | |
| # return [gr.update(interactive=True, value='Train'), gr.update(value='GPU Memory is not enough for training... Please upgrade to GPU atleast 40GB or clone the repo to your local machine.'), None, gr.update()] | |
| if self.training: | |
| return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()] | |
| attributes = attributes_input | |
| if is_person: | |
| attributes = 'white, black, asian, hispanic, indian, male, female' | |
| self.training = True | |
| train_xl(target=target_concept, positive=positive_prompt, negative=negative_prompt, lr=lr_input, iterations=iterations_input, config_file='trainscripts/textsliders/data/config-xl.yaml', rank=int(rank), train_method=train_method_input, device=self.device, attributes=attributes, save_name=save_name) | |
| self.training = False | |
| torch.cuda.empty_cache() | |
| model_map[save_name.replace('.pt','')] = 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'), f'models/{save_name}', gr.update(choices=list(model_map.keys()), value=save_name.replace('.pt',''))] | |
| def inference(self, prompt, seed, start_noise, scale, model_name, model, pbar = gr.Progress(track_tqdm=True)): | |
| seed = seed or 42753 | |
| if self.current_model != model: | |
| if model=='SDXL Turbo': | |
| model_id = "stabilityai/sdxl-turbo" | |
| euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") | |
| self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=euler_anc, torch_dtype=self.weight_dtype).to(self.device) | |
| self.pipe.enable_xformers_memory_efficient_attention() | |
| self.guidance_scale = 1 | |
| self.num_inference_steps = 3 | |
| self.current_model = 'SDXL Turbo' | |
| else: | |
| model_id = 'stabilityai/stable-diffusion-xl-base-1.0' | |
| self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=self.weight_dtype).to(self.device) | |
| self.pipe.enable_xformers_memory_efficient_attention() | |
| self.guidance_scale = 7.5 | |
| self.num_inference_steps = 20 | |
| self.current_model = 'SDXL' | |
| 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(float(model_path.split('_')[-1].replace('.pt',''))) | |
| 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=self.num_inference_steps, generator=generator, network=network, start_noise=int(start_noise), scale=float(scale), unet=unet, guidance_scale=self.guidance_scale).images[0] | |
| generator = torch.manual_seed(seed) | |
| original_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=self.num_inference_steps, generator=generator, network=network, start_noise=start_noise, scale=0, unet=unet, guidance_scale=self.guidance_scale).images[0] | |
| del unet, network | |
| unet = None | |
| network = None | |
| torch.cuda.empty_cache() | |
| return edited_image, original_image | |
| demo = Demo() | |