ZERO GPU

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import spaces
 from PIL import Image, ImageDraw, ImageFont
 from src.condition import Condition
 from diffusers.pipelines import FluxPipeline
@@ -7,22 +8,22 @@ import numpy as np
 
 from src.generate import seed_everything, generate
 
-pipe = None
+# pipe = None
 
 
-def init_pipeline():
-    global pipe
-    pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
-    )
-    pipe = pipe.to("cuda")
-    pipe.load_lora_weights(
-        "Yuanshi/OminiControl",
-        weight_name=f"omini/subject_512.safetensors",
-        adapter_name="subject",
-    )
-
+# def init_pipeline():
+#     global pipe
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
+)
+pipe = pipe.to("cuda")
+pipe.load_lora_weights(
+    "Yuanshi/OminiControl",
+    weight_name=f"omini/subject_512.safetensors",
+    adapter_name="subject",
+)
 
+@spaces.GPU
 def process_image_and_text(image, text):
     # center crop image
     w, h, min_size = image.size[0], image.size[1], min(image.size)
@@ -38,8 +39,8 @@ def process_image_and_text(image, text):
 
     condition = Condition("subject", image)
 
-    if pipe is None:
-        init_pipeline()
+    # if pipe is None:
+    #     init_pipeline()
 
     result_img = generate(
         pipe,
@@ -87,7 +88,7 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    init_pipeline()
+    # init_pipeline()
     demo.launch(
         debug=True,
         ssr_mode=False