update

app.py

@@ -4,18 +4,20 @@ from PIL import Image, ImageDraw, ImageFont
 from src.condition import Condition
 from diffusers.pipelines import FluxPipeline
 import numpy as np
+import devicetorch
 
 from src.generate import seed_everything, generate
 
 pipe = None
 
+device = devicetorch.get(torch)
 
 def init_pipeline():
     global pipe
     pipe = FluxPipeline.from_pretrained(
         "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
     )
-    pipe = pipe.to("cuda")
+    pipe = pipe.to(device)
     pipe.load_lora_weights(
         "Yuanshi/OminiControl",
         weight_name=f"omini/subject_512.safetensors",

requirements.txt

@@ -3,4 +3,5 @@ diffusers
 peft
 opencv-python
 protobuf
-sentencepiece
+sentencepiece
+gradio