Update app.py

app.py

@@ -749,18 +749,6 @@ model = Flux().to(dtype=torch.bfloat16, device="cuda")
 result = model.load_state_dict(sd)
 model_zero_init = False
 
-@spaces.GPU()
-def apply_nf4_storage():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch_device = torch.device(device)
-    
-    global model, model_zero_init
-    if not model_zero_init:
-        model = model.to(torch_device)
-        model_zero_init = True
-
-apply_nf4_storage()
-
 # model = Flux().to(dtype=torch.bfloat16, device="cuda")
 # result = model.load_state_dict(load_file("/storage/dev/nyanko/flux-dev/flux1-dev.sft"))
 
@@ -777,6 +765,11 @@ def generate_image(
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch_device = torch.device(device)
     
+    global model, model_zero_init
+    if not model_zero_init:
+        model = model.to(torch_device)
+        model_zero_init = True
+    
     if do_img2img and init_image is not None:
         init_image = get_image(init_image)
         if resize_img: