Replace DALLE-mini with Stable Diffusion pipeline.

app.py

@@ -12,9 +12,10 @@ from fastai.callback.core import Callback
 from fastai.learner import *
 from fastai.torch_core import TitledStr
 from html2image import Html2Image
-from min_dalle import MinDalle
+# from min_dalle import MinDalle
 from torch import tensor, Tensor, float16, float32
 from torch.distributions import Transform
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 
 # These utility functions need to be in main (or otherwise where created) because fastai loads from that module, see:
 # https://docs.fast.ai/learner.html#load_learner
@@ -50,34 +51,41 @@ class DropOutput(Callback):
 
 # initialize only once
 # Takes about 2 minutes (126 seconds) to generate an image in Huggingface spaces on CPU
-model = MinDalle(
-    models_root='./pretrained',
-    dtype=float32,
-    device='cpu',
-    is_mega=True,
-    is_reusable=True
-)
+# NOTE as of 2022-11-13 min-dalle is broken, switch to using a stable diffusion model for images
+# model = MinDalle(
+#     models_root='./pretrained',
+#     dtype=float32,
+#     device='cpu',
+#     is_mega=True,
+#     is_reusable=True
+# )
+pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", custom_pipeline="stable_diffusion_mega", torch_dtype=torch.float16, revision="fp16")
+
+# pipeline.to("cuda")
 
 
 def gen_image(prompt):
+    prompt = f"{prompt}, fantasy painting by Greg Rutkowski"
     # See https://huggingface.co/spaces/pootow/min-dalle/blob/main/app.py
     # Hugging Space faces seems to run out of memory if grads are not disabled
-    torch.set_grad_enabled(False)
+    # torch.set_grad_enabled(False)
     print(f'RUNNING gen_image with prompt: {prompt}')
-    images = model.generate_images(
-        text=prompt,
-        seed=-1,
-        grid_size=1,  # grid size above 2 causes out of memory on 12 GB 3080Ti; grid size 2 gives 4 images
-        is_seamless=False,
-        temperature=1,
-        top_k=256,
-        supercondition_factor=16,
-        is_verbose=True
-    )
+    images = pipeline.text2img(prompt, width=512, height=512).images
+    # images = model.generate_images(
+    #     text=prompt,
+    #     seed=-1,
+    #     grid_size=1,  # grid size above 2 causes out of memory on 12 GB 3080Ti; grid size 2 gives 4 images
+    #     is_seamless=False,
+    #     temperature=1,
+    #     top_k=256,
+    #     supercondition_factor=16,
+    #     is_verbose=True
+    # )
     print('COMPLETED GENERATION')
-    images = images.to('cpu').numpy()
-    images = images.astype(np.uint8)
-    return Image.fromarray(images[0])
+    # images = images.to('cpu').numpy()
+    # images = images.astype(np.uint8)
+    # return Image.fromarray(images[0])
+    return images[0]
 
 
 gpu = False
@@ -326,7 +334,8 @@ x = gr.components.Textbox()
 iface = gr.Interface(title="MonsterGen", theme="default", description=app_description, fn=run, inputs=[input_box],
                      outputs=[output_monster_card, output_text_box, output_monster_image, output_monster_html])
 iface.launch()
-# TODO: Add examples
+# TODO: Add examples, larger language model?, document process, log silences, "Passives" => "Traits", log timestamps
+# Fine tune dalle-mini? https://blog.paperspace.com/dalle-mini/
 # API works, assuming query takes no longer than 30 seconds (504 gateway timeout)
 # Looks like API page improvements are in progress: https://github.com/gradio-app/gradio/issues/1325
 # Example code below: