import gradio as gr import torch from diffusers import FluxPipeline, StableDiffusion3Pipeline from PIL import Image from typing import Optional import os import random import numpy as np import spaces import huggingface_hub import copy from FlowEdit_utils import FlowEditSD3, FlowEditFLUX SD3STRING = 'stabilityai/stable-diffusion-3-medium-diffusers' FLUXSTRING = 'black-forest-labs/FLUX.1-dev' device = "cuda" if torch.cuda.is_available() else "cpu" # device = "cpu" # model_type = 'SD3' pipe_sd3 = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK')) pipe_flux = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK')) # pipe_sd3.to(device) # pipe_flux.to(device) # scheduler = pipe.scheduler # pipe = pipe.to(device) loaded_model = 'None' def on_model_change(model_type): if model_type == 'SD3': T_steps_value = 50 src_guidance_scale_value = 3.5 tar_guidance_scale_value = 13.5 n_max_value = 33 elif model_type == 'FLUX': T_steps_value = 28 src_guidance_scale_value = 1.5 tar_guidance_scale_value = 5.5 n_max_value = 24 else: raise NotImplementedError(f"Model type {model_type} not implemented") return T_steps_value, src_guidance_scale_value, tar_guidance_scale_value, n_max_value def get_examples(): case = [ ["inputs/cat.png", "SD3", 50, 3.5, 13.5, 33, "a cat sitting in the grass", "a puppy sitting in the grass", 0, 1, 42], ["inputs/iguana.png", "SD3", 50, 3.5, 13.5, 31, "A large orange lizard sitting on a rock near the ocean. The lizard is positioned in the center of the scene, with the ocean waves visible in the background. The rock is located close to the water, providing a picturesque setting for the lizard''s resting spot.", "A large dragon sitting on a rock near the ocean. The dragon is positioned in the center of the scene, with the ocean waves visible in the background. The rock is located close to the water, providing a picturesque setting for the dragon''s resting spot.", 0, 1, 42], ["inputs/cat.png", "FLUX", 28, 1.5, 5.5, 24, "a cat sitting in the grass", "a puppy sitting in the grass", 0, 1, 42], ["inputs/gas_station.png", "FLUX", 28, 1.5, 5.5, 23, "A gas station with a white and red sign that reads \"CAFE\" There are several cars parked in front of the gas station, including a white car and a van.", "A gas station with a white and red sign that reads \"CVPR\" There are several cars parked in front of the gas station, including a white car and a van.", 0, 1, 42], ["inputs/steak.png", "FLUX", 28, 1.5, 5.5, 24, "A steak accompanied by a side of leaf salad.", "A bread roll accompanied by a side of leaf salad.", 0, 1, 42], ["inputs/kill_bill.png", "FLUX", 28, 2.5, 6.5, 22, "a blonde woman in a yellow jumpsuit holding a sword in front of her face", "a blonde woman in a yellow jumpsuit holding a sword in front of her face, anime style drawing", 14, 1, 42], ] return case @spaces.GPU(duration=60) def FlowEditRun( image_src: str, model_type: str, T_steps: int, src_guidance_scale: float, tar_guidance_scale: float, n_max: int, src_prompt: str, tar_prompt: str, n_min: int, n_avg: int, seed: int, ): if not len(src_prompt): raise gr.Error("source prompt cannot be empty") if not len(tar_prompt): raise gr.Error("target prompt cannot be empty") # global pipe_sd3 # global scheduler # global loaded_model # reload model only if different from the loaded model # if loaded_model != model_type: if model_type == 'FLUX': # pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK')) pipe = pipe_flux.to(device) elif model_type == 'SD3': # pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK')) pipe = pipe_sd3.to(device) else: raise NotImplementedError(f"Model type {model_type} not implemented") scheduler = pipe.scheduler # pipe = pipe.to(device) # set seed random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) # load image image = Image.open(image_src) # crop image to have both dimensions divisibe by 16 - avoids issues with resizing image = image.crop((0, 0, image.width - image.width % 16, image.height - image.height % 16)) image_src = pipe.image_processor.preprocess(image) # image_tar = pipe.image_processor.postprocess(image_src) # return image_tar[0] # cast image to half precision image_src = image_src.to(device).half() with torch.autocast("cuda"), torch.inference_mode(): x0_src_denorm = pipe.vae.encode(image_src).latent_dist.mode() x0_src = (x0_src_denorm - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor # send to cuda x0_src = x0_src.to(device) negative_prompt = "" # optionally add support for negative prompts (SD3) if model_type == 'SD3': x0_tar = FlowEditSD3(pipe, scheduler, x0_src, src_prompt, tar_prompt, negative_prompt, T_steps, n_avg, src_guidance_scale, tar_guidance_scale, n_min, n_max,) elif model_type == 'FLUX': x0_tar = FlowEditFLUX(pipe, scheduler, x0_src, src_prompt, tar_prompt, negative_prompt, T_steps, n_avg, src_guidance_scale, tar_guidance_scale, n_min, n_max,) else: raise NotImplementedError(f"Sampler type {model_type} not implemented") x0_tar_denorm = (x0_tar / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor with torch.autocast("cuda"), torch.inference_mode(): image_tar = pipe.vae.decode(x0_tar_denorm, return_dict=False)[0] image_tar = pipe.image_processor.postprocess(image_tar) return image_tar[0] # title = "FlowEdit: Inversion-Free Text-Based Editing Using Pre-Trained Flow Models" intro = """