import os import re import time from io import BytesIO import uuid from dataclasses import dataclass from glob import iglob import argparse from einops import rearrange from fire import Fire from PIL import ExifTags, Image import spaces import torch import torch.nn.functional as F import gradio as gr import numpy as np from transformers import pipeline from flux.sampling import denoise, get_schedule, prepare, unpack from flux.util import (configs, embed_watermark, load_ae, load_clip, load_flow_model, load_t5) from huggingface_hub import login login(token=os.getenv('Token')) import torch device = torch.cuda.current_device() total_memory = torch.cuda.get_device_properties(device).total_memory allocated_memory = torch.cuda.memory_allocated(device) reserved_memory = torch.cuda.memory_reserved(device) print(f"Total memory: {total_memory / 1024**2:.2f} MB") print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB") print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB") @dataclass class SamplingOptions: source_prompt: str target_prompt: str # prompt: str width: int height: int num_steps: int guidance: float seed: int | None offload = False name = "flux-dev" is_schnell = False feature_path = 'feature' output_dir = 'result' add_sampling_metadata = True # class FluxEditor: # def __init__(self, args): # self.args = args # self.device = torch.device(args.device) # self.offload = args.offload # self.name = args.name # self.is_schnell = args.name == "flux-schnell" # self.feature_path = 'feature' # self.output_dir = 'result' # self.add_sampling_metadata = True # if self.name not in configs: # available = ", ".join(configs.keys()) # raise ValueError(f"Got unknown model name: {name}, chose from {available}") # # init all components # if self.offload: # self.model.cpu() # torch.cuda.empty_cache() # self.ae.encoder.to(self.device) ae = load_ae(name, device="cpu" if offload else device) t5 = load_t5(device, max_length=256 if name == "flux-schnell" else 512) clip = load_clip(device) model = load_flow_model(name, device="cpu" if offload else device) print("!!!!!!!!!!!!device!!!!!!!!!!!!!!",device) print("!!!!!!!!self.t5!!!!!!",next(t5.parameters()).device) print("!!!!!!!!self.clip!!!!!!",next(clip.parameters()).device) print("!!!!!!!!self.model!!!!!!",next(model.parameters()).device) @torch.inference_mode() def encode(init_image, torch_device, ae): init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1 init_image = init_image.unsqueeze(0) init_image = init_image.to(torch_device) ae = ae.cuda() with torch.no_grad(): init_image = ae.encode(init_image.to()).to(torch.bfloat16) return init_image @spaces.GPU(duration=120) @torch.inference_mode() def edit(init_image, source_prompt, target_prompt, num_steps, inject_step, guidance, seed): device = "cuda" if torch.cuda.is_available() else "cpu" torch.cuda.empty_cache() seed = None # if seed == -1: # seed = None shape = init_image.shape new_h = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16 new_w = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16 init_image = init_image[:new_h, :new_w, :] width, height = init_image.shape[0], init_image.shape[1] init_image = encode(init_image, device, ae) print(init_image.shape) rng = torch.Generator(device="cpu") opts = SamplingOptions( source_prompt=source_prompt, target_prompt=target_prompt, width=width, height=height, num_steps=num_steps, guidance=guidance, seed=seed, ) if opts.seed is None: opts.seed = torch.Generator(device="cpu").seed() print(f"Generating with seed {opts.seed}:\n{opts.source_prompt}") t0 = time.perf_counter() opts.seed = None #############inverse####################### info = {} info['feature'] = {} info['inject_step'] = inject_step print("!!!!!!!!!!!!device!!!!!!!!!!!!!!",device) print("!!!!!!!!self.t5!!!!!!",next(t5.parameters()).device) print("!!!!!!!!self.clip!!!!!!",next(clip.parameters()).device) print("!!!!!!!!self.model!!!!!!",next(model.parameters()).device) device = torch.cuda.current_device() total_memory = torch.cuda.get_device_properties(device).total_memory allocated_memory = torch.cuda.memory_allocated(device) reserved_memory = torch.cuda.memory_reserved(device) print(f"Total memory: {total_memory / 1024**2:.2f} MB") print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB") print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB") with torch.no_grad(): inp = prepare(t5, clip, init_image, prompt=opts.source_prompt) inp_target = prepare(t5, clip, init_image, prompt=opts.target_prompt) timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell")) # inversion initial noise with torch.no_grad(): z, info = denoise(model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info) inp_target["img"] = z timesteps = get_schedule(opts.num_steps, inp_target["img"].shape[1], shift=(name != "flux-schnell")) # denoise initial noise x, _ = denoise(model, **inp_target, timesteps=timesteps, guidance=guidance, inverse=False, info=info) # decode latents to pixel space x = unpack(x.float(), opts.width, opts.height) output_name = os.path.join(output_dir, "img_{idx}.jpg") if not os.path.exists(output_dir): os.makedirs(output_dir) idx = 0 else: fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)] if len(fns) > 0: idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1 else: idx = 0 ae = ae.cuda() with torch.autocast(device_type=device.type, dtype=torch.bfloat16): x = ae.decode(x) if torch.cuda.is_available(): torch.cuda.synchronize() t1 = time.perf_counter() fn = output_name.format(idx=idx) print(f"Done in {t1 - t0:.1f}s. Saving {fn}") # bring into PIL format and save x = x.clamp(-1, 1) x = embed_watermark(x.float()) x = rearrange(x[0], "c h w -> h w c") img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy()) exif_data = Image.Exif() exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux" exif_data[ExifTags.Base.Make] = "Black Forest Labs" exif_data[ExifTags.Base.Model] = name if add_sampling_metadata: exif_data[ExifTags.Base.ImageDescription] = source_prompt img.save(fn, exif=exif_data, quality=95, subsampling=0) print("End Edit") return img def create_demo(model_name: str, device: str = "cuda:0" if torch.cuda.is_available() else "cpu", offload: bool = False): is_schnell = model_name == "flux-schnell" with gr.Blocks() as demo: gr.Markdown(f"# RF-Edit Demo (FLUX for image editing)") with gr.Row(): with gr.Column(): # source_prompt = gr.Textbox(label="Source Prompt", value="") # target_prompt = gr.Textbox(label="Target Prompt", value="") source_prompt = gr.Text( label="Source Prompt", show_label=False, max_lines=1, placeholder="Enter your source prompt", container=False, value="" ) target_prompt = gr.Text( label="Target Prompt", show_label=False, max_lines=1, placeholder="Enter your target prompt", container=False, value="" ) init_image = gr.Image(label="Input Image", visible=True) generate_btn = gr.Button("Generate") with gr.Column(): with gr.Accordion("Advanced Options", open=True): num_steps = gr.Slider(1, 30, 25, step=1, label="Number of steps") inject_step = gr.Slider(1, 15, 5, step=1, label="Number of inject steps") guidance = gr.Slider(1.0, 10.0, 2, step=0.1, label="Guidance", interactive=not is_schnell) # seed = gr.Textbox(0, label="Seed (-1 for random)", visible=False) # add_sampling_metadata = gr.Checkbox(label="Add sampling parameters to metadata?", value=False) output_image = gr.Image(label="Generated Image") generate_btn.click( fn=edit, inputs=[init_image, source_prompt, target_prompt, num_steps, inject_step, guidance], outputs=[output_image] ) return demo # if __name__ == "__main__": # import argparse # parser = argparse.ArgumentParser(description="Flux") # parser.add_argument("--name", type=str, default="flux-dev", choices=list(configs.keys()), help="Model name") # parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use") # parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use") # parser.add_argument("--share", action="store_true", help="Create a public link to your demo") # parser.add_argument("--port", type=int, default=41035) # args = parser.parse_args() demo = create_demo("flux-dev", "cuda") demo.launch()