fix zerogpu bugs

.gitignore

@@ -1,3 +1,4 @@
 **/.eggs/
 *.so
-JarvisIR/package/agent_tools/RIDCP/.eggs
+JarvisIR/package/agent_tools/RIDCP/.eggs
+JarvisIR/checkpoints

JarvisIR/package/agent_tools/Retinexformer/basicsr_retinexformer/version.py

@@ -1,5 +1,5 @@
 # GENERATED VERSION FILE
-# TIME: Tue Jun 10 10:29:45 2025
+# TIME: Tue Jun 10 11:41:53 2025
 __version__ = '1.2.0+unknown'
 short_version = '1.2.0'
 version_info = (1, 2, 0)

JarvisIR/package/agent_tools/img2img_turbo/inference.py

@@ -15,7 +15,7 @@ def load_turbo_model(name, model_path, device):
     model.direction = 'b2a'
     model.caption = 'driving in the day'
     model.eval()
-    model.unet.enable_xformers_memory_efficient_attention()
+    # model.unet.enable_xformers_memory_efficient_attention()
     return model
 
 

app.py

@@ -7,37 +7,60 @@ import torch
 from PIL import Image
 from transformers import AutoProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
 from threading import Thread
-from agent_tools import RestorationToolkit
+import subprocess
 
-# download tools ckpts
-from huggingface_hub import hf_hub_download, list_files_info
+def install_cuda_toolkit():
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
+    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
 
-# 目标目录
-target_dir = "JarvisIR/checkpoints/agent_tools"
-os.makedirs(target_dir, exist_ok=True)
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
+        os.environ["CUDA_HOME"],
+        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
+    )
+    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
 
-# 获取文件列表
-files = list_files_info(repo_id="LYL1015/JarvisIR", repo_type="model")
+install_cuda_toolkit()
 
-for file in files:
-    if file.path.startswith("agent_tools/checkpoints/"):
-        rel_path = os.path.relpath(file.path, "agent_tools/checkpoints")
-        local_path = os.path.join(target_dir, rel_path)
-        
-        os.makedirs(os.path.dirname(local_path), exist_ok=True)
-        
-        hf_hub_download(
-            repo_id="LYL1015/JarvisIR",
-            filename=file.path,
-            local_dir=target_dir,
-            local_dir_use_symlinks=False,
-            force_filename=local_path  
-        )
+def download_tools_ckpts(target_dir, url):
+    from huggingface_hub import snapshot_download
+    import os
+    import shutil
+
+    tmp_dir = "hf_temp_download"
+    os.makedirs(tmp_dir, exist_ok=True)
+
+    snapshot_download(
+        repo_id="LYL1015/JarvisIR",
+        repo_type="model",
+        local_dir=tmp_dir,
+        allow_patterns=os.path.join(url, "**"), 
+        local_dir_use_symlinks=False,
+    )
+
+    src_dir = os.path.join(tmp_dir, url)
+    
+ 
+    shutil.copytree(src_dir, target_dir) 
+
+    shutil.rmtree(tmp_dir)
 
+target_dir = "JarvisIR/checkpoints/agent_tools"
+if not os.path.exists(target_dir):
+    download_tools_ckpts(target_dir, "agent_tools/checkpoints")
+
+llm_targer_dir = "JarvisIR/checkpoints/pretrained_preview"
+if not os.path.exists(llm_targer_dir):
+    download_tools_ckpts(llm_targer_dir, "pretrained/preview")
 
 # Model configuration
 # XXX: Path to the fine-tuned LLaVA model
-model_id = "LYL1015/JarvisIR" 
+model_id = llm_targer_dir 
     
 # Available image restoration tasks and their corresponding models
 all_tasks = " {denoise: [scunet, restormer], lighten: [retinexformer_fivek, hvicidnet, lightdiff], \
@@ -73,13 +96,13 @@ prompts_query2 = [
 print("Loading LLM model...")
 
 # Initialize the image restoration toolkit
+from agent_tools import RestorationToolkit
 tool_engine = RestorationToolkit(score_weight=[0,0,0,0,0])
 # Load the LLaVA model in half precision to reduce memory usage
 model = LlavaForConditionalGeneration.from_pretrained(
     model_id,
-    subfolder="pretrained/preview",  # 关键参数：指定子目录
     torch_dtype=torch.float16,
-    device_map="auto",  # 比 .to(0) 更智能的设备分配
+    device_map="auto",  
     low_cpu_mem_usage=True
 )
 processor = AutoProcessor.from_pretrained(model_id)
@@ -198,7 +221,6 @@ def resize_image_to_original(processed_image_path, original_size):
         return output_path
     return processed_image_path
 
-@spaces.GPU(duration=150)
 def get_llm_response_streaming(image_path):
     """
     Get streaming response from LLM for image analysis
@@ -236,7 +258,7 @@ def get_llm_response_streaming(image_path):
     
     return streamer
 
-@spaces.GPU(duration=150)
+
 def process_image_with_tools(image_path, models, original_size):
     """
     Process image using the tool engine and restore to original size
@@ -262,7 +284,7 @@ def process_image_with_tools(image_path, models, original_size):
     final_result = resize_image_to_original(res['output_path'], original_size)
     
     return final_result
-
+@spaces.GPU(duration=150)
 def process_full_pipeline(image):
     """
     Main processing pipeline with streaming UI updates