Update

README.md

@@ -4,7 +4,7 @@ emoji: 🏢
 colorFrom: blue
 colorTo: pink
 sdk: gradio
-sdk_version: 3.36.1
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 suggested_hardware: t4-small

app.py

@@ -3,11 +3,23 @@
 from __future__ import annotations
 
 import os
+import shlex
+import subprocess
+import sys
 
 import gradio as gr
+import PIL.Image
+import spaces
 import torch
+from diffusers import DPMSolverMultistepScheduler
 
-from model import Model
+if os.getenv("SYSTEM") == "spaces":
+    with open("patch") as f:
+        subprocess.run(shlex.split("patch -p1"), cwd="multires_textual_inversion", stdin=f)
+
+sys.path.insert(0, "multires_textual_inversion")
+
+from pipeline import MultiResPipeline, load_learned_concepts
 
 DESCRIPTION = "# [Multiresolution Textual Inversion](https://github.com/giannisdaras/multires_textual_inversion)"
 
@@ -20,9 +32,39 @@ For this demo, only `<jane>`, `<gta5-artwork>` and `<cat-toy>` are available.
 Also, `number` should be an integer in [0, 9].
 """
 
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model_id = "runwayml/stable-diffusion-v1-5"
+if device.type == "cpu":
+    pipe = MultiResPipeline.from_pretrained(model_id)
+else:
+    pipe = MultiResPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16")
+pipe = pipe.to(device)
+pipe.scheduler = DPMSolverMultistepScheduler(
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    num_train_timesteps=1000,
+    trained_betas=None,
+    prediction_type="epsilon",
+    thresholding=False,
+    algorithm_type="dpmsolver++",
+    solver_type="midpoint",
+    lower_order_final=True,
+)
+string_to_param_dict = load_learned_concepts(pipe, "textual_inversion_outputs/")
+
+
+@spaces.GPU
+def run(prompt: str, n_images: int, n_steps: int, seed: int) -> list[PIL.Image.Image]:
+    generator = torch.Generator(device=device).manual_seed(seed)
+    return pipe(
+        [prompt] * n_images,
+        string_to_param_dict,
+        num_inference_steps=n_steps,
+        generator=generator,
+    )
 
-model = Model()
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
@@ -40,18 +82,24 @@ with gr.Blocks(css="style.css") as demo:
                     value=1,
                 )
             with gr.Row():
-                num_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=10)
+                num_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=10,
+                )
             with gr.Row():
                 seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, value=100)
             with gr.Row():
-                run_button = gr.Button("Run")
+                run_button = gr.Button()
 
         with gr.Column():
             result = gr.Gallery(label="Result", object_fit="scale-down")
 
     with gr.Row():
         with gr.Group():
-            fn = lambda x: model.run(x, 2, 10, 100)
+            fn = lambda x: run(x, 2, 10, 100)
             with gr.Row():
                 gr.Examples(
                     label="Examples 1",
@@ -64,7 +112,6 @@ with gr.Blocks(css="style.css") as demo:
                     inputs=prompt,
                     outputs=result,
                     fn=fn,
-                    cache_examples=CACHE_EXAMPLES,
                 )
             with gr.Row():
                 gr.Examples(
@@ -78,7 +125,6 @@ with gr.Blocks(css="style.css") as demo:
                     inputs=prompt,
                     outputs=result,
                     fn=fn,
-                    cache_examples=CACHE_EXAMPLES,
                 )
             with gr.Row():
                 gr.Examples(
@@ -91,7 +137,6 @@ with gr.Blocks(css="style.css") as demo:
                     inputs=prompt,
                     outputs=result,
                     fn=fn,
-                    cache_examples=CACHE_EXAMPLES,
                 )
 
         inputs = [
@@ -101,13 +146,13 @@ with gr.Blocks(css="style.css") as demo:
             seed,
         ]
         prompt.submit(
-            fn=model.run,
+            fn=run,
             inputs=inputs,
             outputs=result,
             api_name=False,
         )
         run_button.click(
-            fn=model.run,
+            fn=run,
             inputs=inputs,
             outputs=result,
             api_name="run",
@@ -116,4 +161,5 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Accordion("About available prompts", open=False):
         gr.Markdown(DETAILS)
 
-demo.queue(max_size=10).launch()
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

model.py

@@ -1,48 +0,0 @@
-from __future__ import annotations
-
-import os
-import shlex
-import subprocess
-import sys
-
-import PIL.Image
-import torch
-from diffusers import DPMSolverMultistepScheduler
-
-if os.getenv("SYSTEM") == "spaces":
-    with open("patch") as f:
-        subprocess.run(shlex.split("patch -p1"), cwd="multires_textual_inversion", stdin=f)
-
-sys.path.insert(0, "multires_textual_inversion")
-
-from pipeline import MultiResPipeline, load_learned_concepts
-
-
-class Model:
-    def __init__(self):
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        model_id = "runwayml/stable-diffusion-v1-5"
-        if self.device.type == "cpu":
-            pipe = MultiResPipeline.from_pretrained(model_id)
-        else:
-            pipe = MultiResPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16")
-        self.pipe = pipe.to(self.device)
-        self.pipe.scheduler = DPMSolverMultistepScheduler(
-            beta_start=0.00085,
-            beta_end=0.012,
-            beta_schedule="scaled_linear",
-            num_train_timesteps=1000,
-            trained_betas=None,
-            predict_epsilon=True,
-            thresholding=False,
-            algorithm_type="dpmsolver++",
-            solver_type="midpoint",
-            lower_order_final=True,
-        )
-        self.string_to_param_dict = load_learned_concepts(self.pipe, "textual_inversion_outputs/")
-
-    def run(self, prompt: str, n_images: int, n_steps: int, seed: int) -> list[PIL.Image.Image]:
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-        return self.pipe(
-            [prompt] * n_images, self.string_to_param_dict, num_inference_steps=n_steps, generator=generator
-        )

requirements.txt

@@ -1,6 +1,6 @@
-accelerate==0.20.3
-diffusers==0.9.0
-ftfy==6.1.1
-Pillow==10.0.0
-torch==1.13.1
-transformers==4.30.2
+accelerate==0.31.0
+diffusers==0.28.2
+ftfy==6.2.0
+Pillow==10.3.0
+torch==2.0.1
+transformers==4.41.2