fix pycuda bugs

app.py

@@ -4,21 +4,26 @@ import sys
 import datetime
 import gradio as gr
 import numpy as np
-import random
+from PIL import Image
 import spaces #[uncomment to use ZeroGPU]
-# from diffusers import DiffusionPipeline
 import torch
 from torchvision.transforms import ToTensor, ToPILImage
-import logging
-# logging.getLogger("huggingface_hub").setLevel(logging.CRITICAL)
-from huggingface_hub import hf_hub_download, snapshot_download
 
+
+# -------------------------- HuggingFace -------------------------------
+from huggingface_hub import hf_hub_download, snapshot_download
 model_name = "iimmortall/UltraFusion"
 auth_token = os.getenv("HF_AUTH_TOKEN")
-# greet_file = hf_hub_download(repo_id=model_name, filename="main.py", use_auth_token=auth_token)  
+greet_file = hf_hub_download(repo_id=model_name, filename="main.py", use_auth_token=auth_token)
 model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app")
 
 from ultrafusion_utils import load_model, run_ultrafusion, check_input
+PYCUDA_FLAG = True
+try :
+    import pycuda
+except Exception:
+    PYCUDA_FLAG = False
+    print("No pycuda!!!")
 
 RUN_TIMES = 0
 
@@ -43,14 +48,19 @@ def infer(
 ):
     print(under_expo_img.size)
     print("reciving image")
-    # print(under_expo_img.orig_name, over_expo_img.orig_name)
 
-    # under_expo_img = under_expo_img.resize([1500, 1000])
-    # over_expo_img = over_expo_img.resize([1500, 1000])
-    under_expo_img, over_expo_img = check_input(under_expo_img, over_expo_img, max_l=1500)
+    under_expo_img_lr, over_expo_img_lr, under_expo_img, over_expo_img, use_bgu = check_input(under_expo_img, over_expo_img, max_l=1500)
+
+    global PYCUDA_FLAG
+    if not PYCUDA_FLAG and use_bgu:
+        print("No pycuda, do not run BGU.")
+        use_bgu = False
+
+    ue = to_tensor(under_expo_img_lr).unsqueeze(dim=0).to("cuda")
+    oe = to_tensor(over_expo_img_lr).unsqueeze(dim=0).to("cuda")
+    ue_hr = to_tensor(under_expo_img).unsqueeze(dim=0).to("cuda")
+    oe_hr = to_tensor(over_expo_img).unsqueeze(dim=0).to("cuda")
 
-    ue = to_tensor(under_expo_img).unsqueeze(dim=0).to("cuda")
-    oe = to_tensor(over_expo_img).unsqueeze(dim=0).to("cuda")
     print("num_inference_steps:", num_inference_steps)
     try:
         if num_inference_steps is None:
@@ -59,8 +69,7 @@ def infer(
     except Exception as e:
         num_inference_steps = 20
 
-    out = run_ultrafusion(ue, oe, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, 
-                          steps=num_inference_steps, consistent_start=None)
+    out = run_ultrafusion(ue, oe, ue_hr, oe_hr, use_bgu, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, steps=num_inference_steps, consistent_start=None, test_bs=16)
 
     out = out.clamp(0, 1).squeeze()
     out_pil = to_pil(out)

requirements.txt

@@ -2,14 +2,11 @@ accelerate
 diffusers
 invisible_watermark
 transformers
-xformers
-torch
-torchvision==0.19.1
 omegaconf
 numpy
 pillow
 einops
 scipy
-numpy
 ftfy
-pytorch_lightning==2.4
+pytorch_lightning==2.4
+# pycuda