update

OmniGen/pipeline.py

@@ -18,7 +18,7 @@ from diffusers.utils import (
 )
 
 from OmniGen import OmniGen, OmniGenProcessor, OmniGenScheduler
-
+import devicetorch
 
 logger = logging.get_logger(__name__) 
 
@@ -52,7 +52,7 @@ class OmniGenPipeline:
         self.model = model
         self.processor = processor
 
-        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.device = devicetorch.get(torch)
         self.model.to(self.device)
         self.model.eval()
         self.vae.to(self.device)
@@ -226,4 +226,4 @@ class OmniGenPipeline:
         for i, sample in enumerate(output_samples):  
             output_images.append(Image.fromarray(sample))
         
-        return output_images
+        return output_images

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from PIL import Image
 import os
-import spaces
 
 from OmniGen import OmniGenPipeline
 
@@ -9,7 +8,6 @@ pipe = OmniGenPipeline.from_pretrained(
     "Shitao/OmniGen-v1"
 )
 
-@spaces.GPU(duration=180)
 # 示例处理函数：生成图像
 def generate_image(text, img1, img2, img3, height, width, guidance_scale, inference_steps, seed):
     input_images = [img1, img2, img3]
@@ -250,4 +248,4 @@ with gr.Blocks() as demo:
     )
 
 # 启动应用
-demo.launch()
+demo.launch()

requirements.txt

@@ -1,9 +1,7 @@
 accelerate
 diffusers
 invisible_watermark
-torch
 transformers
-xformers
 timm
 peft
-safetensors
+safetensors