added upsampling

app.py

@@ -3,12 +3,14 @@ import numpy as np
 import torch as th
 from imageio import imread
 from skimage.transform import resize as imresize
+from PIL import Image
 
-from ema_pytorch import EMA
 from decomp_diffusion.model_and_diffusion_util import *
 from decomp_diffusion.diffusion.respace import SpacedDiffusion
 from decomp_diffusion.gen_image import *
+
 from download import download_model
+from upsampling import get_pipeline, upscale_image
 
 import gradio as gr
 
@@ -149,17 +151,19 @@ def combine_components_slice(model, gd, im1, im2, indices=None, sample_method='d
     return sample[0].cpu()
 
 
-
 def decompose_image_demo(im, model):
     sample_method = 'ddim'
     result = gen_image_and_components(MODELS[model], GD[sample_method], im, sample_method=sample_method, num_images=1, device=device)
-    return result.permute(1, 2, 0).numpy()
+    result = Image.fromarray(result.permute(1, 2, 0).numpy())
+    return result
 
 
 def combine_images_demo(im1, im2, model):
     sample_method = 'ddim'
     result = combine_components_slice(MODELS[model], GD[sample_method], im1, im2, indices='1,0,1,0', sample_method=sample_method, num_images=1, device=device)
-    return result.permute(1, 2, 0).numpy()
+    result = Image.fromarray(result.permute(1, 2, 0).numpy())
+    if model == 'CelebA-HQ':
+        return upscale_image(result, pipe)
 
 
 def load_model(dataset, extra_kwargs={}, device='cuda'):
@@ -189,6 +193,7 @@ MODELS = {
     'CelebA-HQ': celeb_model
 }
 
+pipe = get_pipeline()
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -221,7 +226,7 @@ with gr.Blocks() as demo:
                 )
             
         with gr.Column():
-            decomp_output = gr.Image(type='numpy')
+            decomp_output = gr.Image(type='pil')
             decomp_button = gr.Button("Generate")
     
 
@@ -253,7 +258,7 @@ with gr.Blocks() as demo:
             
 
         with gr.Column(scale=1):
-            comb_output = gr.Image(type='numpy')
+            comb_output = gr.Image(type='pil')
             comb_button = gr.Button("Generate")
     
 

requirements.txt

@@ -1,2 +1,5 @@
 git+https://github.com/jsu27/decomp_diffusion.git
 scikit-image==0.19.2
+diffusers==0.10.2
+transformers
+pillow

upsampling.py

@@ -0,0 +1,43 @@
+import torch as th
+from diffusers import IFImg2ImgSuperResolutionPipeline
+from transformers import T5EncoderModel
+from PIL import Image 
+import numpy as np
+
+def get_pipeline():
+
+    text_encoder = T5EncoderModel.from_pretrained(
+        "DeepFloyd/IF-I-XL-v1.0",
+        subfolder="text_encoder", 
+        device_map="auto", 
+        load_in_8bit=True, 
+        variant="8bit"
+    )
+
+    pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained(
+        "DeepFloyd/IF-II-L-v1.0", 
+        text_encoder=text_encoder, 
+        variant="fp16", 
+        torch_dtype=th.float16, 
+        device_map="auto",
+        watermarker=None
+    )
+    return pipe
+
+
+
+def upscale_image(im, pipe):
+    """im is 64x64 PIL image"""
+    prompt = ''
+    prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
+    generator = th.Generator().manual_seed(0)
+
+    image = pipe(
+        image=original_image,
+        original_image=original_image,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_embeds, 
+        generator=generator,
+    ).images[0]
+
+    return image