modified: app.py

app.py

@@ -41,6 +41,11 @@ 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)
+t5.eval()
+clip.eval()
+ae.eval()
+model.eval()
+
 is_schnell = False
 add_sampling_metadata = True
 
@@ -51,8 +56,14 @@ if os.path.exists("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}'
+exp_folders = [d for d in os.listdir(out_root) if d.startswith("exp_") and d[4:].isdigit()]
+if exp_folders:
+    max_idx = max(int(d[4:]) for d in exp_folders)
+    name_dir = f"exp_{max_idx + 1}"
+else:
+    name_dir = "exp_0"
 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"):
@@ -73,11 +84,17 @@ def reset():
     if os.path.exists("history_gradio/history.safetensors"):
         os.remove("history_gradio/history.safetensors")
 
-    out_root = 'src/gradio_utils/gradio_outputs'
+    global out_root, output_dir, history_tensors, source_image, instructions
     if not os.path.exists(out_root):
         os.makedirs(out_root)
-    name_dir = f'exp_{len(os.listdir(out_root))+1}'
+    exp_folders = [d for d in os.listdir(out_root) if d.startswith("exp_") and d[4:].isdigit()]
+    if exp_folders:
+        max_idx = max(int(d[4:]) for d in exp_folders)
+        name_dir = f"exp_{max_idx + 1}"
+    else:
+        name_dir = "exp_0"
     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"):
@@ -128,7 +145,7 @@ def generate_image(
     init_image=None,
     image2image_strength=0.0,
 ):
-    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload
+    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload, history_tensors
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.cuda.empty_cache()
     seed = None
@@ -251,12 +268,6 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
     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
@@ -333,7 +344,6 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
         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)
@@ -460,6 +470,7 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
 
 def on_select(gallery, selected: gr.SelectData):
     return gallery[selected.index][0], gallery[selected.index][1]
+    #return gallery[selected.index][0]
 
 def on_upload(path, uploaded: gr.EventData):
     return path[0][0]
@@ -474,7 +485,7 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
     
     # Pre-defined examples
     examples = [
-        ["src/gradio_utils/gradio_examples/000000000011.jpg", "", "a photo of a eagle standing on the branch", ['attn_guidance'], 15, 3.5, 11, 0],
+        ["src/gradio_utils/gradio_examples/000000000011.jpg", "", "an eagle standing on the branch", ['attn_guidance'], 15, 3.5, 11, 0],
         ["src/gradio_utils/gradio_examples/221000000002.jpg", "", "a cat wearing a hat standing on the fence", ['attn_guidance'], 15, 3.5, 11, 0],
     ]
 
@@ -511,6 +522,7 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
                 example_image = gr.Image(label="example Image", visible=False, width=200)
 
         gallery.select(on_select, gallery, [init_image, source_prompt])
+        #gallery.select(on_select, gallery, [init_image])
         gallery.upload(on_upload, gallery, init_image)
         example_image.change(on_change, example_image, [gallery, init_image])
 
@@ -542,4 +554,4 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
 
 demo = create_demo(name, "cuda")
 #demo.launch(server_name='0.0.0.0', share=args.share, server_port=args.port)
-demo.launch(debug=True)
+demo.launch(debug=True)