Update app.py

GPT4: It seems like your app is exceeding the available memory limit of 46GB. As there are many models being loaded, one possible solution is to load models on-demand when required instead of loading all of them from the beginning. You can do this by changing the `models` list into a dictionary and modifying the init method of the `Model` class to load the models only when specifically requested.

app.py

@@ -22,18 +22,8 @@ class Model:
     def __init__(self, name, path=""):
         self.name = name
         self.path = path
-
-        if path != "":
-            self.pipe_t2i = StableDiffusionPipeline.from_pretrained(
-                path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER
-            )
-            self.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config(
-                self.pipe_t2i.scheduler.config
-            )
-            self.pipe_i2i = StableDiffusionImg2ImgPipeline(**self.pipe_t2i.components)
-        else:
-            self.pipe_t2i = None
-            self.pipe_i2i = None
+        self.pipe_t2i = None
+        self.pipe_i2i = None
 
 
 models = [
@@ -51,6 +41,19 @@ MODELS = {m.name: m for m in models}
 
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
+def get_model(name):
+    moel = MODELS[name]
+
+    if model.pipe_t2i is None:
+        model.pip_t2i = StableDiffusionPipeline.from_pretrained(
+            model.path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER
+        )
+        model.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config(
+            mode.pipe_t2i.scheduler.config
+        )
+        model.pipe_i2i = StableDiffusionImg2ImgPipeline(**model.pipe_t2i.components)
+
+        return model
 
 def error_str(error, title="Error"):
     return (
@@ -170,7 +173,7 @@ def img_to_img(
     generator,
     seed,
 ):
-    pipe = MODELS[model_name].pipe_i2i
+    pipe = model.pipe_i2i
 
     if torch.cuda.is_available():
         pipe = pipe.to("cuda")