modified: app.py

app.py

@@ -32,441 +32,430 @@ class SamplingOptions:
     guidance: float
     seed: int | None
 
-@torch.inference_mode()
-def encode(init_image, torch_device, ae):
-    if next(ae.parameters()).device != torch_device:
-        ae = ae.to(torch_device)
-    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)
-    with torch.no_grad():
-        init_image = ae.encode(init_image).to(torch.bfloat16)
-    return init_image
-
 
-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"
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+offload = False
+device = "cuda" if torch.cuda.is_available() else "cpu"
+name = 'flux-dev'
+ae = load_ae(name, device="cpu" if offload else torch_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 torch_device)
+is_schnell = False
+add_sampling_metadata = True
+
+# clear history
+if os.path.exists("history_gradio/history.safetensors"):
+    os.remove("history_gradio/history.safetensors")
+
+out_root = 'src/gradio_utils/gradio_outputs'
+if not os.path.exists(out_root):
+    os.makedirs(out_root)
+name_dir = f'exp_{len(os.listdir(out_root))+1}'
+output_dir = os.path.join(out_root, name_dir)
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+if not os.path.exists("heatmap"):
+    os.makedirs("heatmap")
+if not os.path.exists("heatmap/average_heatmaps"):
+    os.makedirs("heatmap/average_heatmaps")
+source_image = None
+history_tensors = {
+        "source img": torch.zeros((1, 1, 1)),
+        "prev img": torch.zeros((1, 1, 1))}   
+instructions = ['source']
 
-        self.feature_path = 'feature'
 
-        self.reset()
-        
-        self.add_sampling_metadata = True
-
-        if self.name not in configs:
-            available = ", ".join(configs.keys())
-            raise ValueError(f"Got unknown model name: {self.name}, chose from {available}")
-
-        # init all components
-        self.clip = load_clip(self.device)
-        self.t5 = load_t5(self.device, max_length=256 if self.name == "flux-schnell" else 77)
-        self.model = load_flow_model(self.name, self.device)
-        self.ae = load_ae(self.name, self.device)
-        self.t5.eval()
-        self.clip.eval()
-        self.ae.eval()
-        self.model.eval()
-
-        # clear history
-        if os.path.exists("history_gradio/history.safetensors"):
-            os.remove("history_gradio/history.safetensors")
+@torch.inference_mode()
+def reset():
+
+    # clear history
+    if os.path.exists("history_gradio/history.safetensors"):
+        os.remove("history_gradio/history.safetensors")
+
+    out_root = 'src/gradio_utils/gradio_outputs'
+    if not os.path.exists(out_root):
+        os.makedirs(out_root)
+    name_dir = f'exp_{len(os.listdir(out_root))+1}'
+    output_dir = os.path.join(out_root, name_dir)
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    if not os.path.exists("heatmap"):
+        os.makedirs("heatmap")
+    if not os.path.exists("heatmap/average_heatmaps"):
+        os.makedirs("heatmap/average_heatmaps")
+    instructions = ['source']
+    source_image = None
+    history_tensors = {
+        "source img": torch.zeros((1, 1, 1)),
+        "prev img": torch.zeros((1, 1, 1))}   
     
+    source_prompt = "(Optional) Describe the content of the uploaded image."
+    traget_prompt = "(Required) Describe the desired content of the edited image."
+    gallery = None
+    output_image = None
+    return source_prompt, traget_prompt, gallery, output_image
 
-    @torch.inference_mode()
-    def reset(self):
-        out_root = 'src/gradio_utils/gradio_outputs'
-        if not os.path.exists(out_root):
-            os.makedirs(out_root)
-        name_dir = f'exp_{len(os.listdir(out_root))}'
-        self.output_dir = os.path.join(out_root, name_dir)
-        if not os.path.exists(self.output_dir):
-            os.makedirs(self.output_dir)
-        if not os.path.exists("heatmap"):
-            os.makedirs("heatmap")
-        if not os.path.exists("heatmap/average_heatmaps"):
-            os.makedirs("heatmap/average_heatmaps")
-        self.instructions = ['source']
-        self.source_image = None
-        self.history_tensors = {
-            "source img": torch.zeros((1, 1, 1)),
-            "prev img": torch.zeros((1, 1, 1))}   
-        
-        source_prompt = "(Optional) Describe the content of the uploaded image."
-        traget_prompt = "(Required) Describe the desired content of the edited image."
-        gallery = None
-        output_image = None
-        return source_prompt, traget_prompt, gallery, output_image
-
-
-    @torch.inference_mode()
-    def process_image(self, 
-                    init_image,
-                    source_prompt, 
-                    target_prompt, 
-                    editing_strategy, 
-                    denoise_strategy, 
-                    num_steps, 
-                    guidance, 
-                    attn_guidance_start_block, 
-                    inject_step, 
-                    init_image_2=None):
-        if init_image is None:
-            img, gr_gallery = self.generate_image(prompt=target_prompt)
-        else:
-            img, gr_gallery = self.edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2)
-        return img, gr_gallery
+
+@torch.inference_mode()
+def process_image(
+                init_image,
+                source_prompt, 
+                target_prompt, 
+                editing_strategy, 
+                denoise_strategy, 
+                num_steps, 
+                guidance, 
+                attn_guidance_start_block, 
+                inject_step, 
+                init_image_2=None):
+    if init_image is None:
+        img, gr_gallery = generate_image(prompt=target_prompt)
+    else:
+        img, gr_gallery = edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2)
+    return img, gr_gallery
 
         
-    @spaces.GPU(duration=120)
-    @torch.inference_mode()
-    def generate_image(
-        self,
-        width=512,
-        height=512,
-        num_steps=28,
-        guidance=3.5,
-        seed=None,
-        prompt='',
-        init_image=None,
-        image2image_strength=0.0,
-        add_sampling_metadata=True,
-    ):
-
-        if seed is None:
-            g_seed = torch.Generator(device=torch.device("cpu")).seed()
-        print(f"Generating '{prompt}' with seed {g_seed}")
-        t0 = time.perf_counter()
-
-        if init_image is not None:
-            if isinstance(init_image, np.ndarray):
-                init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 255.0
-                init_image = init_image.unsqueeze(0)
-            init_image = init_image.to(self.device)
-            init_image = torch.nn.functional.interpolate(init_image, (height, width))
-            if self.offload:
-                self.ae.encoder.to(self.device)
-            init_image = self.ae.encode(init_image)
-            if self.offload:
-                self.ae = self.ae.cpu()
-                torch.cuda.empty_cache()
-
-        # prepare input
-        x = get_noise(
-            1,
-            height,
-            width,
-            device=self.device,
-            dtype=torch.bfloat16,
-            seed=g_seed,
-        )
-        timesteps = get_schedule(
-            num_steps,
-            x.shape[-1] * x.shape[-2] // 4,
-            shift=(not self.is_schnell),
-        )
-        if init_image is not None:
-            t_idx = int((1 - image2image_strength) * num_steps)
-            t = timesteps[t_idx]
-            timesteps = timesteps[t_idx:]
-            x = t * x + (1.0 - t) * init_image.to(x.dtype)
-
-        if self.offload:
-            self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
-        inp = prepare(t5=self.t5, clip=self.clip, img=x, prompt=prompt)
-
-        # offload TEs to CPU, load model to gpu
-        if self.offload:
-            self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
-            torch.cuda.empty_cache()
-            self.model = self.model.to(self.device)
-
-        self.model = self.model.to(self.device)
-        # denoise initial noise
-        info = {}
-        info['feature'] = {}
-        info['inject_step'] = 0
-        info['editing_strategy']= ""
-        info['start_layer_index'] = 0
-        info['end_layer_index'] = 37
-        info['reuse_v']= False
-        qkv_ratio = '1.0,1.0,1.0'
-        info['qkv_ratio'] = list(map(float, qkv_ratio.split(',')))
-        x = denoise_rf(self.model, **inp, timesteps=timesteps, guidance=guidance, inverse=False, info=info)
-
-        # offload model, load autoencoder to gpu
-        if self.offload:
-            self.model.cpu()
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def generate_image(
+    width=512,
+    height=512,
+    num_steps=28,
+    guidance=3.5,
+    seed=None,
+    prompt='',
+    init_image=None,
+    image2image_strength=0.0,
+):
+    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.cuda.empty_cache()
+    seed = None
+
+    if seed is None:
+        g_seed = torch.Generator(device="cpu").seed()
+    print(f"Generating '{prompt}' with seed {g_seed}")
+    t0 = time.perf_counter()
+
+    if init_image is not None:
+        if isinstance(init_image, np.ndarray):
+            init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 255.0
+            init_image = init_image.unsqueeze(0)
+        init_image = init_image.to(device)
+        init_image = torch.nn.functional.interpolate(init_image, (height, width))
+        if offload:
+            ae.encoder.to(device)
+        init_image = ae.encode(init_image)
+        if offload:
+            ae = ae.cpu()
             torch.cuda.empty_cache()
-            self.ae.decoder.to(x.device)
 
-        # decode latents to pixel space
-        x = unpack(x[0].float(), height, width)
-        self.ae = self.ae.to(x.device)
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
-            x = self.ae.decode(x)
+    # prepare input
+    x = get_noise(
+        1,
+        height,
+        width,
+        device=device,
+        dtype=torch.bfloat16,
+        seed=g_seed,
+    )
+    timesteps = get_schedule(
+        num_steps,
+        x.shape[-1] * x.shape[-2] // 4,
+        shift=(not is_schnell),
+    )
+    if init_image is not None:
+        t_idx = int((1 - image2image_strength) * num_steps)
+        t = timesteps[t_idx]
+        timesteps = timesteps[t_idx:]
+        x = t * x + (1.0 - t) * init_image.to(x.dtype)
+
+    if offload:
+        t5, clip = t5.to(device), clip.to(device)
+    inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
+
+    # offload TEs to CPU, load model to gpu
+    if offload:
+        t5, clip = t5.cpu(), clip.cpu()
+        torch.cuda.empty_cache()
+        model = model.to(device)
+
+    # denoise initial noise
+    info = {}
+    info['feature'] = {}
+    info['inject_step'] = 0
+    info['editing_strategy']= ""
+    info['start_layer_index'] = 0
+    info['end_layer_index'] = 37
+    info['reuse_v']= False
+    qkv_ratio = '1.0,1.0,1.0'
+    info['qkv_ratio'] = list(map(float, qkv_ratio.split(',')))
+    x = denoise_rf(model, **inp, timesteps=timesteps, guidance=guidance, inverse=False, info=info)
+
+    # offload model, load autoencoder to gpu
+    if offload:
+        model.cpu()
+        torch.cuda.empty_cache()
+        ae.decoder.to(x.device)
 
-        if self.offload:
-            self.ae.decoder.cpu()
-            torch.cuda.empty_cache()
+    # decode latents to pixel space
+    x = unpack(x[0].float(), height, width)
+    device = torch.device("cuda")
+    with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+        x = ae.decode(x)
 
-        t1 = time.perf_counter()
+    if offload:
+        ae.decoder.cpu()
+        torch.cuda.empty_cache()
 
-        print(f"Done in {t1 - t0:.1f}s.")
-        # bring into PIL format
-        x = x.clamp(-1, 1)
-        x = embed_watermark(x.float())
-        x = rearrange(x[0], "c h w -> h w c")
+    t1 = time.perf_counter()
 
-        img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
-        
-        filename = os.path.join(self.output_dir,f"round_0000_[{prompt}].jpg")
-        os.makedirs(os.path.dirname(filename), exist_ok=True)
-        exif_data = Image.Exif()
-        if init_image is None:
-            exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
-        else:
-            exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
-        exif_data[ExifTags.Base.Make] = "Black Forest Labs"
-        exif_data[ExifTags.Base.Model] = self.name
-        if add_sampling_metadata:
-            exif_data[ExifTags.Base.ImageDescription] = prompt
-        img.save(filename, format="jpeg", exif=exif_data, quality=95, subsampling=0)
-        self.instructions = [prompt]
-
-            #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
-        img_and_prompt = []
-        history_imgs = sorted(os.listdir(self.output_dir))
-        for img_file, prompt_txt in zip(history_imgs, self.instructions):
-            img_and_prompt.append((os.path.join(self.output_dir, img_file), prompt_txt))
-            history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
-        return img, history_gallery
-
-
-    @spaces.GPU(duration=200)
-    @torch.inference_mode()
-    def edit(self, init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2=None):
-        
-        torch.cuda.empty_cache()
-        seed = None
+    print(f"Done in {t1 - t0:.1f}s.")
+    # bring into PIL format
+    x = x.clamp(-1, 1)
+    x = embed_watermark(x.float())
+    x = rearrange(x[0], "c h w -> h w c")
 
-        print(f"Inital_t5_device: {self.t5.hf_module.device}")
-        print(f"Inital_clip_device: {self.clip.hf_module.device}")
-        print(f"Inital_flow_model: {self.model.img_in.weight.device}")
-        print(f"Inital_flow_model self.model.img_in: {self.model.img_in}")
-        print(f"Inital_flow_model self.model.time_in.out_layer.weight: {self.model.time_in.out_layer.weight}")
-        
-        if self.offload:
-            self.model.cpu()
-            torch.cuda.empty_cache()
-            self.ae.encoder.to(self.device)
-
-        #-----------------------------     0.1 prepare multi-turn editing     -------------------------------------#
-        info = {}
-        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
-
-        if not any("round_0000" in fname for fname in os.listdir(self.output_dir)):
-            Image.fromarray(init_image).save(os.path.join(self.output_dir,"round_0000_[source].jpg"))
-
-
-        init_image = init_image[:new_h, :new_w, :]
-        width, height = init_image.shape[0], init_image.shape[1]
-        init_image = encode(init_image, self.device, self.ae)
-
-        print(init_image.shape)
-
-        if init_image_2 is None:
-            print("init_image_2 is not provided, proceeding with single image processing.")
-        else: 
-            init_image_2_pil = Image.fromarray(init_image_2) # Convert NumPy array to PIL Image
-            init_image_2_pil = init_image_2_pil.resize((new_w, new_h), Image.Resampling.LANCZOS) 
-            init_image_2 = np.array(init_image_2_pil)  # Convert back to NumPy (if needed)
-            init_image_2 = encode(init_image_2, self.device, self.ae)
-
-        rng = torch.Generator(device=torch.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=torch.device("cpu")).seed()
-        
-        print(f"Editing with prompt:\n{opts.source_prompt}")
-        t0 = time.perf_counter()
+    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+    
+    filename = os.path.join(output_dir,f"round_0000_[{prompt}].jpg")
+    os.makedirs(os.path.dirname(filename), exist_ok=True)
+    exif_data = Image.Exif()
+    if init_image is None:
+        exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
+    else:
+        exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
+    exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+    exif_data[ExifTags.Base.Model] = name
+    if add_sampling_metadata:
+        exif_data[ExifTags.Base.ImageDescription] = prompt
+    img.save(filename, format="jpeg", exif=exif_data, quality=95, subsampling=0)
+    instructions = [prompt]
 
-        opts.seed = None
-        if self.offload:
-            self.ae = self.ae.cpu()
-            torch.cuda.empty_cache()
-            self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
-
-        #-----------------------------     0.2 prepare attention strategy     -------------------------------------#
-        info = {}
-        info['feature'] = {}
-        info['inject_step'] = inject_step
-        info['editing_strategy']= " ".join(editing_strategy)
-        info['start_layer_index'] = 0
-        info['end_layer_index'] = 37
-        info['reuse_v']= False
-        qkv_ratio = '1.0,1.0,1.0'
-        info['qkv_ratio'] = list(map(float, qkv_ratio.split(',')))
-        info['attn_guidance'] = attn_guidance_start_block
-        info['lqr_stop'] = 0.25
-
-        if not os.path.exists(self.feature_path):
-            os.mkdir(self.feature_path)
-
-
-        #-----------------------------     0.3 prepare latents     -------------------------------------#
-        with torch.no_grad():
-            inp = prepare(self.t5, self.clip, init_image, prompt=opts.source_prompt)
-            inp_target = prepare(self.t5, self.clip, init_image, prompt=opts.target_prompt)
-            if self.source_image is None:
-                self.source_image = inp['img']
-            inp_target_2 = None
-            if not init_image_2 is None:
-                inp_target_2 = prepare_image(init_image_2)
-                info['lqr_stop'] = 0.35
-
-        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
-        #timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=False)
-
-        # offload TEs to CPU, load model to gpu
-        if self.offload:
-            self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
-            torch.cuda.empty_cache()
-            self.model = self.model.to(self.device)
+        #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
+    img_and_prompt = []
+    history_imgs = sorted(os.listdir(output_dir))
+    for img_file, prompt_txt in zip(history_imgs, instructions):
+        img_and_prompt.append((os.path.join(output_dir, img_file), prompt_txt))
+        history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
+    return img, history_gallery
 
-        self.model = self.model.to(self.device)
-        print(f"model has been moved to {self.device}")
-        
-        #-----------------------------     1 Inverting current image     -------------------------------------#
-        denoise_strategies = ['fireflow', 'rf', 'rf_solver', 'midpoint', 'rf_inversion', 'multi_turn_consistent']
-        denoise_funcs = [denoise_fireflow, denoise_rf, denoise_rf_solver, denoise_midpoint, denoise_rf_inversion, denoise_multi_turn_consistent]
-        denoise_func = denoise_funcs[denoise_strategies.index(denoise_strategy)]
-        with torch.no_grad():
-            z, info = denoise_func(self.model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info)
-        
-        print(f"Inverted Z: {z}")
-        print(info)
-        
-        
-        #-----------------------------     2 history_tensors used to implement dual-LQR guiding editing     -------------------------------------#
-        inp_target["img"] = z
-        timesteps = get_schedule(opts.num_steps, inp_target["img"].shape[1], shift=(self.name != "flux-schnell"))
-
-        if torch.all(self.history_tensors['source img'] == 0):
-            self.history_tensors = {
-            "source img": inp["img"],
-            "prev img": inp_target_2} 
-        else:
-            if inp_target_2 is None:
-                self.history_tensors["prev img"] = inp["img"]
-            else:
-                self.history_tensors["source img"] = inp["img"]
-                self.history_tensors["prev img"] = inp_target_2
-
-        #-----------------------------     3 sampling     -------------------------------------#
-        if denoise_strategy in ['rf_inversion', 'multi_turn_consistent']:
-            x, _ = denoise_func(self.model, **inp_target, timesteps=timesteps, guidance=guidance, inverse=False, info=info, img_LQR=self.history_tensors)
-        else:
-            x, _ = denoise_func(self.model, **inp_target, timesteps=timesteps, guidance=opts.guidance, inverse=False, info=info)
+
+@spaces.GPU(duration=200)
+@torch.inference_mode()
+def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2=None):
+    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload, source_image, history_tensors, instructions
+   
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.cuda.empty_cache()
+    seed = None
+
+    print(f"Inital_t5_device: {t5.hf_module.device}")
+    print(f"Inital_clip_device: {clip.hf_module.device}")
+    print(f"Inital_flow_model: {model.img_in.weight.device}")
+    print(f"Inital_flow_model self.model.img_in: {model.img_in}")
+    print(f"Inital_flow_model self.model.time_in.out_layer.weight: {model.time_in.out_layer.weight}")
+
+    #-----------------------------     0.1 prepare multi-turn editing     -------------------------------------#
+    info = {}
+    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
+
+    if not any("round_0000" in fname for fname in os.listdir(output_dir)):
+        Image.fromarray(init_image).save(os.path.join(output_dir,"round_0000_[source].jpg"))
+
+    init_image = init_image[:new_h, :new_w, :]
+    width, height = init_image.shape[0], init_image.shape[1]
+
+    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(device)
+    if offload:
+        model.cpu()
+        torch.cuda.empty_cache()
+        ae.encoder.to(device)
         
+    with torch.no_grad():
+        init_image = ae.encode(init_image.to()).to(torch.bfloat16)
+
+    if init_image_2 is None:
+        print("init_image_2 is not provided, proceeding with single image processing.")
+    else: 
+        init_image_2_pil = Image.fromarray(init_image_2) # Convert NumPy array to PIL Image
+        init_image_2_pil = init_image_2_pil.resize((new_w, new_h), Image.Resampling.LANCZOS) 
+        init_image_2 = np.array(init_image_2_pil)  # Convert back to NumPy (if needed)
+        init_image_2 = torch.from_numpy(init_image_2).permute(2, 0, 1).float() / 127.5 - 1
+
+    rng = torch.Generator(device=torch.device("cpu"))
+    opts = SamplingOptions(
+        source_prompt=source_prompt,
+        target_prompt=target_prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=None,
+    )
+    if opts.seed is None:
+        opts.seed = torch.Generator(device=torch.device("cpu")).seed()
+    
+    print(f"Editing with prompt:\n{opts.source_prompt}")
+    t0 = time.perf_counter()
+
+    if offload:
+        ae = ae.cpu()
+        torch.cuda.empty_cache()
+        t5, clip = t5.to(torch_device), clip.to(torch_device)
+    opts.seed = None
+
+
+    #-----------------------------     0.2 prepare attention strategy     -------------------------------------#
+    info = {}
+    info['feature'] = {}
+    info['inject_step'] = inject_step
+    info['editing_strategy']= " ".join(editing_strategy)
+    info['start_layer_index'] = 0
+    info['end_layer_index'] = 37
+    info['reuse_v']= False
+    qkv_ratio = '1.0,1.0,1.0'
+    info['qkv_ratio'] = list(map(float, qkv_ratio.split(',')))
+    info['attn_guidance'] = attn_guidance_start_block
+    info['lqr_stop'] = 0.25
+
+    #-----------------------------     0.3 prepare latents     -------------------------------------#
+    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)
+        if source_image is None:
+            source_image = inp['img']
+        inp_target_2 = None
+        if not init_image_2 is None:
+            inp_target_2 = prepare_image(init_image_2)
+            info['lqr_stop'] = 0.35
+
+    timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
+    #timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=False)
+
+    # offload TEs to CPU, load model to gpu
+    
+    if offload:
+        t5, clip = t5.cpu(), clip.cpu()
+        torch.cuda.empty_cache()
+        model = model.to(torch_device)
 
-        #-----------------------------     4 update history_tensors     -------------------------------------#
-        info = {}
-        self.history_tensors["source img"] = self.source_image
-        self.history_tensors["prev img"] = x
-        '''save_file(history_tensors, "history_gradio/history.safetensors")'''
+    #-----------------------------     1 Inverting current image     -------------------------------------#
+    denoise_strategies = ['fireflow', 'rf', 'rf_solver', 'midpoint', 'rf_inversion', 'multi_turn_consistent']
+    denoise_funcs = [denoise_fireflow, denoise_rf, denoise_rf_solver, denoise_midpoint, denoise_rf_inversion, denoise_multi_turn_consistent]
+    denoise_func = denoise_funcs[denoise_strategies.index(denoise_strategy)]
+    with torch.no_grad():
+        z, info = denoise_func(model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info)
+    
+    
+    #-----------------------------     2 history_tensors used to implement dual-LQR guiding editing     -------------------------------------#
+    inp_target["img"] = z
+    timesteps = get_schedule(opts.num_steps, inp_target["img"].shape[1], shift=(name != "flux-schnell"))
+
+    if torch.all(history_tensors['source img'] == 0):
+        history_tensors = {
+        "source img": inp["img"],
+        "prev img": inp_target_2} 
+    else:
+        if inp_target_2 is None:
+            history_tensors["prev img"] = inp["img"]
+        else:
+            history_tensors["source img"] = inp["img"]
+            history_tensors["prev img"] = inp_target_2
+
+    #-----------------------------     3 sampling     -------------------------------------#
+    if denoise_strategy in ['rf_inversion', 'multi_turn_consistent']:
+        x, _ = denoise_func(model, **inp_target, timesteps=timesteps, guidance=guidance, inverse=False, info=info, img_LQR=history_tensors)
+    else:
+        x, _ = denoise_func(model, **inp_target, timesteps=timesteps, guidance=opts.guidance, inverse=False, info=info)
+    
 
-        # offload model, load autoencoder to gpu
-        if self.offload:
-            self.model.cpu()
-            torch.cuda.empty_cache()
-            self.ae.decoder.to(x.device)
+    #-----------------------------     4 update history_tensors     -------------------------------------#
+    info = {}
+    history_tensors["source img"] = source_image
+    history_tensors["prev img"] = x
 
+    #-----------------------------     5 decode x to image      -------------------------------------#
+    x = unpack(x.float(), opts.width, opts.height)
 
+    if offload:
+        model.cpu()
+        torch.cuda.empty_cache()
+        ae.decoder.to(x.device)
+            
+    device = torch.device("cuda")
+    with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+        x = ae.decode(x)
 
-        #-----------------------------     5 decode x to image      -------------------------------------#
-        x = unpack(x.float(), opts.width, opts.height)
 
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
-            x = self.ae.decode(x)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    t1 = time.perf_counter()
 
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        t1 = time.perf_counter()
+    # 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")
 
-        # 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 = 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] = self.name
-        if self.add_sampling_metadata:
-            exif_data[ExifTags.Base.ImageDescription] = source_prompt
-    
 
 
-        #--------------------------------     6 save image      -------------------------------------#
+    #--------------------------------     6 save image      -------------------------------------#
 
-        #--------------------     6.1 prepare output folder      ----------------------#
-        if not os.path.exists(self.output_dir):
-            os.makedirs(self.output_dir)
-            idx = 0
-        #--------------------    6.2 editing round      ----------------------#           
+    #--------------------     6.1 prepare output folder      ----------------------#
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 1
+    #--------------------    6.2 editing round      ----------------------#           
+    else:
+        fns = [fn for fn in os.listdir(output_dir)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[1]) for fn in fns) + 1
         else:
-            fns = [fn for fn in os.listdir(self.output_dir)]
-            if len(fns) > 0:
-                idx = max(int(fn.split("_")[1]) for fn in fns) + 1
-            else:
-                idx = 0
-        formatted_idx = str(idx).zfill(4) # Format as a 4-digit string
-
-        #--------------------    6.3 output name      ----------------------#
-        if denoise_strategy == 'multi_turn_consistent':
-            denoise_strategy = 'MTC'
-        if target_prompt == '':
-            target_prompt = 'Reconstruction'
-        if target_prompt == source_prompt:
-            target_prompt = 'Reconstruction: ' + target_prompt
-
-        target_suffix = " ".join(target_prompt.split()[-5:])
-        output_name = f"round_{formatted_idx}_{target_suffix}_{denoise_strategy}.jpg"
-        
-        fn = os.path.join(self.output_dir, output_name)
-        
-        print(f"Done in {t1 - t0:.1f}s. Saving {fn}")
-        img.save(fn)
+            idx = 1
+    formatted_idx = str(idx).zfill(4) # Format as a 4-digit string
+
+    #--------------------    6.3 output name      ----------------------#
+    if denoise_strategy == 'multi_turn_consistent':
+        denoise_strategy = 'MTC'
+    if target_prompt == '':
+        target_prompt = 'Reconstruction'
+    if target_prompt == source_prompt:
+        target_prompt = 'Reconstruction: ' + target_prompt
+
+    target_suffix = " ".join(target_prompt.split()[-5:])
+    output_name = f"round_{formatted_idx}_{target_suffix}_{denoise_strategy}.jpg"
+    
+    fn = os.path.join(output_dir, output_name)
+    
+    print(f"Done in {t1 - t0:.1f}s. Saving {fn}")
+    img.save(fn)
 
-        if 'Reconstruction' in target_prompt:
-            target_prompt = source_prompt
-        self.instructions.append(target_prompt)
-        print("End Edit")
+    if 'Reconstruction' in target_prompt:
+        target_prompt = source_prompt
+    instructions.append(target_prompt)
+    print("End Edit")
 
-        #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
-        img_and_prompt = []
-        history_imgs = sorted(os.listdir(self.output_dir))
-        for img_file, prompt_txt in zip(history_imgs, self.instructions):
-            img_and_prompt.append((os.path.join(self.output_dir, img_file), prompt_txt))
-        history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
-    
-        return img, history_gallery
+    #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
+    img_and_prompt = []
+    history_imgs = sorted(os.listdir(output_dir))
+    for img_file, prompt_txt in zip(history_imgs, instructions):
+        img_and_prompt.append((os.path.join(output_dir, img_file), prompt_txt))
+    history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
+
+    return img, history_gallery
 
 
 def on_select(gallery, selected: gr.SelectData):
@@ -480,8 +469,7 @@ def on_change(init_image, changed: gr.EventData):
     return gr.Gallery(value=[(img_path[0], "")], label="History Image", interactive=True, columns=3), img_path[0]
 
 
-def create_demo(model_name: str, device: str | torch.device = "cuda" if torch.cuda.is_available() else "cpu", offload: bool = False):
-    editor = FluxEditor(args)
+def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
     is_schnell = model_name == "flux-schnell"
     
     # Pre-defined examples
@@ -527,11 +515,11 @@ def create_demo(model_name: str, device: str | torch.device = "cuda" if torch.cu
         example_image.change(on_change, example_image, [gallery, init_image])
 
         generate_btn.click(
-            fn=editor.process_image,
+            fn=process_image,
             inputs=[init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2],
             outputs=[output_image, gallery]
         )
-        reset_btn.click(fn = editor.reset, outputs=[source_prompt, target_prompt, gallery, output_image])
+        reset_btn.click(fn = reset, outputs=[source_prompt, target_prompt, gallery, output_image])
         
         # Add examples
         gr.Examples(
@@ -552,17 +540,6 @@ def create_demo(model_name: str, device: str | torch.device = "cuda" if torch.cu
     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", help="Device to use")
-    parser.add_argument("--offload", default=False, 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=9090)
-    args = parser.parse_args()
-    print(vars(args))
-
-    demo = create_demo(args.name, args.device, args.offload)
-    #demo.launch(server_name='0.0.0.0', share=args.share, server_port=args.port)
-    demo.launch(debug=True)
+demo = create_demo(name, "cuda")
+#demo.launch(server_name='0.0.0.0', share=args.share, server_port=args.port)
+demo.launch(debug=True)