Update app.py

app.py

@@ -23,9 +23,7 @@ import string
 import random, time, math   
 import os
 
-
-os.environ["NCCL_P2P_DISABLE"]="1"
-os.environ["NCCL_IB_DISABLE"]="1"
+os.environ["XVERSE_OUTPUT_DATA"] = f"{os.getcwd()}/output_results"
 
 import src.flux.generate
 from src.flux.generate import generate_from_test_sample, seed_everything
@@ -118,7 +116,6 @@ model.config = config
 
 run_mode = "mod_only"
 store_attn_map = False
-run_name = time.strftime("%m%d-%H%M")
 
 num_inputs = 2
 
@@ -257,10 +254,10 @@ def generate_image(
 
     src_inputs = []
     use_words = []
-    cur_run_time = time.strftime("%m%d-%H%M%S")
-    tmp_dir_root = f"tmp/gradio_demo/{run_name}"
-    temp_dir = f"{tmp_dir_root}/{session_id}/{cur_run_time}_{generate_random_string(4)}"
-    os.makedirs(temp_dir, exist_ok=True)
+    
+    temp_dir = os.path.join(os.environ["XVERSE_OUTPUT_DATA"], session_id)
+    os.makedirs(base_dir, exist_ok=True)
+    
     print(f"Temporary directory created: {temp_dir}")
     for i, (image_path, caption) in enumerate(zip(images, captions)):
         if image_path:
@@ -273,7 +270,7 @@ def generate_image(
                 prompt = prompt.replace(f"ENT{i+1}", caption)
             
             image = resize_keep_aspect_ratio(Image.open(image_path), 768)
-            save_path = f"{temp_dir}/{session_id}/tmp_resized_input_{i}.png"
+            save_path = f"{temp_dir}/tmp_resized_input_{i}.png"
             image.save(save_path)
             
             input_image_path = save_path
@@ -311,7 +308,7 @@ def generate_image(
             ),
         ]
     
-    json_dump(test_sample, f"{temp_dir}/{session_id}/test_sample.json", 'utf-8')
+    json_dump(test_sample, f"{temp_dir}/test_sample.json", 'utf-8')
     assert single_attention == True
     target_size = int(round((target_width * target_height) ** 0.5) // 16 * 16)
     print(test_sample)
@@ -453,10 +450,8 @@ def cleanup(request: gr.Request):
     """
     sid = request.session_hash
     if sid:
-        d1 = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], sid)
-        d2 = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], sid)
+        d1 = os.path.join(os.environ["XVERSE_OUTPUT_DATA"], sid)
         shutil.rmtree(d1, ignore_errors=True)
-        shutil.rmtree(d2, ignore_errors=True)
         
 
 css = """
@@ -590,9 +585,9 @@ if __name__ == "__main__":
                 db_latent_lora_scale_str, sb_latent_lora_scale_str, vae_lora_scale_str,
                 indices_state,
                 session_state,
-                *images,  
-                *captions, 
-                *idip_checkboxes,
+                [*images],  
+                [*captions], 
+                [*idip_checkboxes],
             ], 
             outputs=output
         )