OmniGen

Runtime error

yrr commited on Oct 21

Commit

1876385

•

1 Parent(s): 6128ddf

zero gpu

Files changed (4) hide show

app.py CHANGED Viewed

@@ -1,14 +1,17 @@
 import gradio as gr
 from PIL import Image
 import os
-os.environ['CUDA_VISIBLE_DEVICES'] = '7'
 from OmniGen import OmniGenPipeline
 pipe = OmniGenPipeline.from_pretrained("shitao/tmp-preview")
 # 示例处理函数：生成图像
 def generate_image(text, img1, img2, img3, height, width, guidance_scale):
     input_images = [img1, img2, img3]
     # 去除 None

 import gradio as gr
 from PIL import Image
 import os
+import spaces
+# os.environ['CUDA_VISIBLE_DEVICES'] = '7'
 from OmniGen import OmniGenPipeline
 pipe = OmniGenPipeline.from_pretrained("shitao/tmp-preview")
+pipe.to("cuda")
 # 示例处理函数：生成图像
+@spaces.GPU
 def generate_image(text, img1, img2, img3, height, width, guidance_scale):
     input_images = [img1, img2, img3]
     # 去除 None

imgs/test_cases/img1.jpg ADDED Viewed

imgs/test_cases/img2.jpg ADDED Viewed

imgs/test_cases/img3.jpg ADDED Viewed