Update

app.py

@@ -7,16 +7,15 @@ import random
 
 import gradio as gr
 import numpy as np
+import PIL.Image
 import torch
-
-from model import Model
+from diffusers import UniDiffuserPipeline
 
 DESCRIPTION = "# [UniDiffuser](https://github.com/thu-ml/unidiffuser)"
 
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶</p>"
 
-model = Model()
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -27,6 +26,67 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if torch.cuda.is_available():
+    pipe = UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1", torch_dtype=torch.float16)
+    pipe.to(device)
+
+
+def run(
+    mode: str,
+    prompt: str,
+    image: PIL.Image.Image | None,
+    seed: int = 0,
+    num_steps: int = 20,
+    guidance_scale: float = 8.0,
+) -> tuple[PIL.Image.Image | None, str]:
+    generator = torch.Generator(device=device).manual_seed(seed)
+    if mode == "t2i":
+        pipe.set_text_to_image_mode()
+        sample = pipe(prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        return sample.images[0], ""
+    elif mode == "i2t":
+        pipe.set_image_to_text_mode()
+        sample = pipe(image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        return None, sample.text[0]
+    elif mode == "joint":
+        pipe.set_joint_mode()
+        sample = pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        return sample.images[0], sample.text[0]
+    elif mode == "i":
+        pipe.set_image_mode()
+        sample = pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        return sample.images[0], ""
+    elif mode == "t":
+        pipe.set_text_mode()
+        sample = pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        return None, sample.text[0]
+    elif mode == "i2t2i":
+        pipe.set_image_to_text_mode()
+        sample = pipe(image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        pipe.set_text_to_image_mode()
+        sample = pipe(
+            prompt=sample.text[0],
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        return sample.images[0], ""
+    elif mode == "t2i2t":
+        pipe.set_text_to_image_mode()
+        sample = pipe(prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+        pipe.set_image_to_text_mode()
+        sample = pipe(
+            image=sample.images[0],
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        return None, sample.text[0]
+    else:
+        raise ValueError
+
+
 def create_demo(mode_name: str) -> gr.Blocks:
     with gr.Blocks() as demo:
         with gr.Row():
@@ -82,7 +142,7 @@ def create_demo(mode_name: str) -> gr.Blocks:
             outputs=seed,
             queue=False,
         ).then(
-            fn=model.run,
+            fn=run,
             inputs=[
                 mode,
                 prompt,

model.py

@@ -1,78 +0,0 @@
-from __future__ import annotations
-
-import PIL.Image
-import torch
-from diffusers import UniDiffuserPipeline
-
-
-class Model:
-    def __init__(self):
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        if self.device.type == "cuda":
-            self.pipe = UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1", torch_dtype=torch.float16)
-            self.pipe.to(self.device)
-        else:
-            self.pipe = UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1")
-
-    def run(
-        self,
-        mode: str,
-        prompt: str,
-        image: PIL.Image.Image | None,
-        seed: int = 0,
-        num_steps: int = 20,
-        guidance_scale: float = 8.0,
-    ) -> tuple[PIL.Image.Image | None, str]:
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-        if mode == "t2i":
-            self.pipe.set_text_to_image_mode()
-            sample = self.pipe(
-                prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
-            )
-            return sample.images[0], ""
-        elif mode == "i2t":
-            self.pipe.set_image_to_text_mode()
-            sample = self.pipe(
-                image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
-            )
-            return None, sample.text[0]
-        elif mode == "joint":
-            self.pipe.set_joint_mode()
-            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
-            return sample.images[0], sample.text[0]
-        elif mode == "i":
-            self.pipe.set_image_mode()
-            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
-            return sample.images[0], ""
-        elif mode == "t":
-            self.pipe.set_text_mode()
-            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
-            return None, sample.text[0]
-        elif mode == "i2t2i":
-            self.pipe.set_image_to_text_mode()
-            sample = self.pipe(
-                image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
-            )
-            self.pipe.set_text_to_image_mode()
-            sample = self.pipe(
-                prompt=sample.text[0],
-                num_inference_steps=num_steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-            )
-            return sample.images[0], ""
-        elif mode == "t2i2t":
-            self.pipe.set_text_to_image_mode()
-            sample = self.pipe(
-                prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
-            )
-            self.pipe.set_image_to_text_mode()
-            sample = self.pipe(
-                image=sample.images[0],
-                num_inference_steps=num_steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-            )
-            return None, sample.text[0]
-        else:
-            raise ValueError