Update app.py

app.py

@@ -61,20 +61,23 @@ def run_schnell(prompt):
     print("Ran!")
     return image
 
+@spaces.GPU
 def run_parallel_models(prompt):
-    with ProcessPoolExecutor(max_workers=3) as executor:
-        future_dev_hyper = executor.submit(run_dev_hyper, prompt)
-        future_dev_turbo = executor.submit(run_dev_turbo, prompt)
-        future_schnell = executor.submit(run_schnell, prompt)
-        
-        res_dev_hyper = future_dev_hyper.result()
-        yield res_dev_hyper, gr.update(), gr.update()
-        res_dev_turbo = future_dev_turbo.result()
-        yield gr.update(), res_dev_turbo, gr.update()
-        res_schnell = future_schnell.result()
-        yield gr.update(), gr.update(), res_dev_turbo
-
-run_parallel_models.zerogpu = True
+    pipe_dev.to("cuda")
+    pipe_dev.load_lora_weights(hyper_lora)
+    image = pipe_dev(prompt, num_inference_steps=8, joint_attention_kwargs={"scale": 0.125}).images[0]
+    pipe_dev.unload_lora_weights()
+    yield image, gr.update(), gr.update()
+    pipe_dev.load_lora_weights(turbo_lora)
+    image = pipe_dev(prompt, num_inference_steps=8).images[0]
+    yield gr.update(), image, gr.update()
+    pipe_dev.unload_lora_weights()
+    pipe_dev.to("cpu")
+    pipe_schnell.to("cuda")
+    image = pipe_schnell(prompt, num_inference_steps=4).images[0]
+    yield gr.update(), gr.update(), image
+   
+#run_parallel_models.zerogpu = True
 
 with gr.Blocks() as demo:
     gr.Markdown("# Low Step Flux Comparison")