Update app.py

app.py

@@ -28,14 +28,13 @@ pipe.load_lora_weights(f'{pipeline_path}/pytorch_lora_weights.safetensors')
 pipe.fuse_lora()
 pipe.unload_lora_weights()
 pipe.force_zeros_for_empty_prompt = False
+pipe = EllaXLPipeline(pipe,f'{pipeline_path}/pytorch_model.bin')
 
-def get_pipe():
-    global pipe
-    if type(pipe)!=EllaXLPipeline:
-        pipe.to("cuda")
-        pipe = EllaXLPipeline(pipe,f'{pipeline_path}/pytorch_model.bin')
-    
-    return pipe
+def tocuda():
+    pipe.pipe.vae.to('cuda')
+    pipe.t5_encoder.to('cuda')
+    pipe.pipe.unet.unet.to('cuda')
+    pipe.pipe.unet.ella.to('cuda')
 
 
 # print("Optimizing BRIA-2.4 - this could take a while")
@@ -62,7 +61,8 @@ def get_pipe():
 @spaces.GPU(enable_queue=True)
 def infer(prompt,negative_prompt,seed,resolution, steps):
     
-    pipe = get_pipe()
+    if 'cuda' not in pipe.pipe.device.type:
+        tocuda()
         
     print(f"""
     —/n