import gradio as gr import numpy as np import random import torch import spaces from PIL import Image import os from huggingface_hub import hf_hub_download import torch from diffusers import DiffusionPipeline from huggingface_hub import hf_hub_download #from gradio_imageslider import ImageSlider # Constants MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" import numpy as np MULTIMODAL_VITAL_LAYERS = [0, 1, 17, 18] SINGLE_MODAL_VITAL_LAYERS = list(np.array([28, 53, 54, 56, 25]) - 19) pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")) pipe.fuse_lora(lora_scale=0.125) #pipe.enable_lora() pipe.to(DEVICE, dtype=torch.float16) def get_examples(): case = [ [Image.open("metal.png"),"dragon.png", "a dragon, in 3d melting gold metal",0.9, 0.5, 0, 5, 28, 28, 0, False,False, 2, False, "text/image guided stylzation" ], [Image.open("doll.png"),"anime.png", "anime illustration",0.9, 0.5, 0, 6, 28, 28, 0, False, False, 2, False,"text/image guided stylzation" ], [Image.open("doll.png"), "raccoon.png", "raccoon, made of yarn",0.9, 0.5, 0, 4, 28, 28, 0, False, False, 2, False, "local subject edits" ], [Image.open("cat.jpg"),"parrot.png", "a parrot", 0.9 ,0.5,2, 8,28, 28,0, False , False, 1, False, "local subject edits"], [Image.open("cat.jpg"),"tiger.png", "a tiger", 0.9 ,0.5,0, 4,8, 8,789385745, False , False, 1, True, "local subject edits"], [Image.open("metal.png"), "dragon.png","a dragon, in 3d melting gold metal",0.9, 0.5, 0, 4, 8, 8, 789385745, False,True, 2, True , "text/image guided stylzation"], ] return case def reset_image_input(): return True def reset_do_inversion(image_input): if image_input: return True else: return False def resize_img(image, max_size=1024): width, height = image.size scaling_factor = min(max_size / width, max_size / height) new_width = int(width * scaling_factor) new_height = int(height * scaling_factor) return image.resize((new_width, new_height), Image.LANCZOS) @torch.no_grad() @spaces.GPU(duration=85) def image2latent(image, latent_nudging_scalar = 1.15): image = pipe.image_processor.preprocess(image, height=1024, width=1024,).type(pipe.vae.dtype).to("cuda") latents = pipe.vae.encode(image)["latent_dist"].mean latents = (latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor latents = latents * latent_nudging_scalar height = pipe.default_sample_size * pipe.vae_scale_factor width = pipe.default_sample_size * pipe.vae_scale_factor num_channels_latents = pipe.transformer.config.in_channels // 4 height = 2 * (height // (pipe.vae_scale_factor * 2)) width = 2 * (width // (pipe.vae_scale_factor * 2)) latents = pipe._pack_latents( latents=latents, batch_size=1, num_channels_latents=num_channels_latents, height=height, width=width ) return latents def check_hyper_flux_lora(enable_hyper_flux): if enable_hyper_flux: pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), lora_scale=0.125) pipe.fuse_lora(lora_scale=0.125) return 8, 8 else: pipe.unfuse_lora() return 28, 28 def convert_string_to_list(s): return [int(x) for x in s.split(',') if x] @spaces.GPU(duration=150) def invert_and_edit(image, source_prompt, edit_prompt, multimodal_layers, single_layers, num_inversion_steps, num_inference_steps, seed, randomize_seed, latent_nudging_scalar, guidance_scale, width = 1024, height = 1024, inverted_latent_list = None, do_inversion = True, image_input = False, ): if randomize_seed: seed = random.randint(0, MAX_SEED) if image_input and (image is not None): if do_inversion: inverted_latent_list = pipe( source_prompt, height=1024, width=1024, guidance_scale=1, output_type="pil", num_inference_steps=num_inversion_steps, max_sequence_length=512, latents=image2latent(image, latent_nudging_scalar), invert_image=True ) do_inversion = False else: # move to gpu because of zero and gr.states inverted_latent_list = [tensor.to(DEVICE) for tensor in inverted_latent_list] num_inference_steps = num_inversion_steps latents = inverted_latent_list[-1].tile(2, 1, 1) guidance_scale = [1,3] image_input = True else: latents = torch.randn( (4096, 64), generator=torch.Generator(0).manual_seed(0), dtype=torch.float16, device=DEVICE, ).tile(2, 1, 1) guidance_scale = guidance_scale image_input = False try: multimodal_layers = convert_string_to_list(multimodal_layers) single_layers = convert_string_to_list(single_layers) except: multimodal_layers = MULTIMODAL_VITAL_LAYERS single_layers = SINGLE_MODAL_VITAL_LAYERS output = pipe( [source_prompt, edit_prompt], height=1024, width=1024, guidance_scale=guidance_scale, output_type="pil", num_inference_steps=num_inference_steps, max_sequence_length=512, latents=latents, inverted_latent_list=inverted_latent_list, mm_copy_blocks=multimodal_layers, single_copy_blocks=single_layers, ).images # move back to cpu because of zero and gr.states if inverted_latent_list is not None: inverted_latent_list = [tensor.cpu() for tensor in inverted_latent_list] if image is None: image = output[0] return image, output[1], inverted_latent_list, do_inversion, image_input, seed # UI CSS css = """ #col-container { margin: 0 auto; max-width: 960px; } """ # Create the Gradio interface with gr.Blocks(css=css) as demo: inverted_latents = gr.State() do_inversion = gr.State(False) image_input = gr.State(False) with gr.Column(elem_id="col-container"): gr.Markdown(f"""# Stable Flow 🌊🖌️ ### Edit real images with FLUX.1 [dev] following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Free Image Editing* by Avrahami et al.](https://arxiv.org/pdf/2411.14430) [[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[project page](https://omriavrahami.com/stable-flow/) [[arxiv](https://arxiv.org/pdf/2411.14430)] """) with gr.Row(): with gr.Column(): input_image = gr.Image( label="Input Image", type="pil" ) source_prompt = gr.Text( label="Source Prompt", max_lines=1, placeholder="describe the edited output", ) edit_prompt = gr.Text( label="Edit Prompt", max_lines=1, placeholder="describe the edited output", ) with gr.Row(): multimodal_layers = gr.Text( info = "MMDiT attention layers used for editing", label="vital multimodal layers", max_lines=1, value="0, 1, 17, 18", ) single_layers = gr.Text( info = "DiT attention layers used editing", label="vital single layers", max_lines=1, value="9, 34, 35, 37, 6", ) with gr.Row(): enable_hyper_flux = gr.Checkbox(label="8-step LoRA", value=False, info="may reduce edit quality", visible=False) run_button = gr.Button("Edit", variant="primary") with gr.Column(): result = gr.Image(label="Result") # with gr.Column(): # with gr.Group(): # result = ImageSlider(position=0.5) with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): num_inference_steps = gr.Slider( label="num inference steps", minimum=1, maximum=50, step=1, value=8, ) guidance_scale = gr.Slider( label="guidance scale", minimum=1, maximum=25, step=1, value=3.5, ) with gr.Row(): num_inversion_steps = gr.Slider( label="num inversion steps", minimum=1, maximum=50, step=1, value=25, ) latent_nudging_scalar= gr.Slider( label="latent nudging scalar", minimum=1, maximum=5, step=0.01, value=1.15, ) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, ) run_button.click( fn=invert_and_edit, inputs=[ input_image, source_prompt, edit_prompt, multimodal_layers, single_layers, num_inversion_steps, num_inference_steps, seed, randomize_seed, latent_nudging_scalar, guidance_scale, width, height, inverted_latents, do_inversion, image_input ], outputs=[input_image, result, inverted_latents, do_inversion, image_input, seed], ) # gr.Examples( # examples=get_examples(), # inputs=[input_image,result, prompt, num_inversion_steps, num_inference_steps, seed, randomize_seed, enable_hyper_flux ], # outputs=[result], # ) input_image.input(fn=reset_image_input, outputs=[image_input]).then( fn=reset_do_inversion, inputs = [image_input], outputs=[do_inversion] ) source_prompt.change( fn=reset_do_inversion, inputs = [image_input], outputs=[do_inversion] ) num_inversion_steps.change( fn=reset_do_inversion, inputs = [image_input], outputs=[do_inversion] ) seed.change( fn=reset_do_inversion, inputs = [image_input], outputs=[do_inversion] ) enable_hyper_flux.change( fn=check_hyper_flux_lora, inputs=[enable_hyper_flux], outputs=[num_inversion_steps, num_inference_steps] ) if __name__ == "__main__": demo.launch()