fenfan commited on
Commit
3994a93
·
verified ·
1 Parent(s): def2fd8

chore: update app.py for zero gpu

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -16,6 +16,7 @@ import dataclasses
16
 
17
  import gradio as gr
18
  import torch
 
19
 
20
  from uno.flux.pipeline import UNOPipeline
21
 
@@ -26,6 +27,7 @@ def create_demo(
26
  offload: bool = False,
27
  ):
28
  pipeline = UNOPipeline(model_type, device, offload, only_lora=True, lora_rank=512)
 
29
 
30
  with gr.Blocks() as demo:
31
  gr.Markdown(f"# UNO by UNO team")
 
16
 
17
  import gradio as gr
18
  import torch
19
+ import spaces
20
 
21
  from uno.flux.pipeline import UNOPipeline
22
 
 
27
  offload: bool = False,
28
  ):
29
  pipeline = UNOPipeline(model_type, device, offload, only_lora=True, lora_rank=512)
30
+ pipeline.__call__ = spaces.GPU(duratioin=120)(pipeline.__call__)
31
 
32
  with gr.Blocks() as demo:
33
  gr.Markdown(f"# UNO by UNO team")