support super res

.gitattributes

@@ -38,3 +38,9 @@ sd1.5-engine/superx4.plan filter=lfs diff=lfs merge=lfs -text
 sd1.5-engine/unet_fp16.plan filter=lfs diff=lfs merge=lfs -text
 sd1.5-engine/vae-decoder.plan filter=lfs diff=lfs merge=lfs -text
 sd1.5-engine/vae-encoder.plan filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/scheduler filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/superx4-512-512.plan filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/text_encoder filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/tokenizer filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/vae filter=lfs diff=lfs merge=lfs -text
+sd1.5-engine/feature_extractor filter=lfs diff=lfs merge=lfs -text

demo.py

@@ -1,12 +1,12 @@
 from lyraSD import LyraSD
 
 t2imodel = LyraSD("text2img", "./sd1.5-engine")
-t2imodel.inference(prompt="A fantasy landscape, trending on artstation", use_super=False)
+t2imodel.inference(prompt="A fantasy landscape, trending on artstation", use_super=True)
 
 
 from PIL import Image
 i2imodel = LyraSD("img2img", "./sd1.5-engine")
-demo_img = Image.open("output/text2img_demo.jpg")
+demo_img = Image.open("output/img2img_input.jpg")
 i2imodel.inference(prompt="A fantasy landscape, trending on artstation", 
-                   image=demo_img)
+                   image=demo_img, use_super=True)
 

lyraSD/inference.py

@@ -78,7 +78,8 @@ class LyraSD(object):
 
 
         for i in range(result_image.shape[0]):
-            result_image = Image.fromarray(np.uint8(result_image[i]))
+            result_image =  Image.fromarray(result_image[0])
+            result_image = result_image.resize((512, 512))
             result_image.save(os.path.join(save_dir, save_basename + "-{}.jpg".format(i)))
                 
         return result_image

lyraSD/muse_trt/sd_img2img.py

@@ -84,10 +84,8 @@ class TRTStableDiffusionImg2ImgPipeline:
 
         self.super = SuperX4TRTInfer(
             engine_dir,
-            model_name='superx4.plan',
+            model_name='superx4-512-512.plan',
             fp16=True,
-            o_height=o_height,
-            o_width=o_width
         )
 
     def runEngine(self, model_name, feed_dict):
@@ -359,7 +357,12 @@ class TRTStableDiffusionImg2ImgPipeline:
 
         # 10. SuperX4 Resolution, Optional.
         if use_super:
-            image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
+            image = np.ascontiguousarray(np.transpose(image, (0, 3, 1, 2))).astype(np.float16)
+            #image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
+            image = self.super.infer(image)
+            image = np.uint8(np.transpose(image, (0, 2, 3, 1)))
+        else:
+            image = np.uint8(image)
 
         return image
 

lyraSD/muse_trt/sd_text2img.py

@@ -58,11 +58,9 @@ class TRTStableDiffusionText2ImgPipeline:
         self.stream = cuda.Stream()
 
         self.super = SuperX4TRTInfer(
-            engine_dir,
-            model_name='superx4.plan',
-            fp16=True,
-            o_height=o_height,
-            o_width=o_width
+            engine_dir=engine_dir,
+            model_name='superx4-512-512.plan',
+            fp16=True
         )
 
     def runEngine(self, model_name, feed_dict):
@@ -285,6 +283,10 @@ class TRTStableDiffusionText2ImgPipeline:
 
         # 9. SuperX4 Resolution, Optional.
         if use_super:
-            image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
-            image = np.transpose(image, (0,3,1,2))
+            image = np.ascontiguousarray(np.transpose(image, (0, 3, 1, 2))).astype(np.float16)
+            #image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
+            image = self.super.infer(image)
+            image = np.uint8(np.transpose(image, (0, 2, 3, 1)))
+        else:
+            image = np.uint8(image)
         return image

lyraSD/muse_trt/utilities.py

@@ -534,3 +534,5 @@ def save_image(images, image_path_dir, image_name_prefix):
         image_path = os.path.join(image_path_dir, image_name_prefix+str(i+1)+'-'+str(random.randint(1000, 9999))+'.png')
         print(f"Saving image {i+1} / {images.shape[0]} to: {image_path}")
         Image.fromarray(images[i]).save(image_path)
+
+

sd1.5-engine/{superx4.plan → superx4-512-512.plan}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2145679c82ae3abe7a8f15259fb9bbeb88a232a8c5501699ca824f518cdbf130
-size 10344612
+oid sha256:b4c37b1fa1a9966104975398cf7b4b4a3ed722335ccde47145c9eb8316030797
+size 10211800