Update app.py

app.py

@@ -10,7 +10,7 @@ from glide_text2im.model_creation import (
 )
 has_cuda = th.cuda.is_available()
 device = th.device('cpu' if not has_cuda else 'cuda')
-# Create base model.
+# a base model.
 options = model_and_diffusion_defaults()
 options['use_fp16'] = has_cuda
 options['timestep_respacing'] = '100' # use 100 diffusion steps for fast sampling
@@ -21,7 +21,8 @@ if has_cuda:
 model.to(device)
 model.load_state_dict(load_checkpoint('base', device))
 print('total base parameters', sum(x.numel() for x in model.parameters()))
-# Create upsampler model.
+
+# Create an upsampler model.
 options_up = model_and_diffusion_defaults_upsampler()
 options_up['use_fp16'] = has_cuda
 options_up['timestep_respacing'] = 'fast27' # use 27 diffusion steps for very fast sampling
@@ -38,31 +39,15 @@ def show_images(batch: th.Tensor):
     reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])
     display(Image.fromarray(reshaped.numpy()))
 # Sampling parameters
-prompt = "an oil painting of a corgi"
+prompt = ""
 batch_size = 1
 guidance_scale = 3.0
 
 # Tune this parameter to control the sharpness of 256x256 images.
 # A value of 1.0 is sharper, but sometimes results in grainy artifacts.
 upsample_temp = 0.997
-
-samples = diffusion.p_sample_loop(
-    model_fn,
-    (full_batch_size, 3, options["image_size"], options["image_size"]),
-    device=device,
-    clip_denoised=True,
-    progress=True,
-    model_kwargs=model_kwargs,
-    cond_fn=None,
-)[:batch_size]
-model.del_cache()
-
-# Show the output
-show_images(samples)
-
-
 import gradio as gr
-def generate_upsampled_image_from_text(prompt):
+def generate_image_from_text(prompt):
     # Set the prompt text
     prompt = prompt
 
@@ -71,43 +56,54 @@ def generate_upsampled_image_from_text(prompt):
     ##############################
 
     # Create the text tokens to feed to the model.
-    tokens = model_up.tokenizer.encode(prompt)
-    tokens, mask = model_up.tokenizer.padded_tokens_and_mask(
-        tokens, options_up['text_ctx']
+    tokens = model.tokenizer.encode(prompt)
+    tokens, mask = model.tokenizer.padded_tokens_and_mask(
+        tokens, options['text_ctx']
     )
 
-    # Create the model conditioning dict.
-    model_kwargs = dict(
-        # Low-res image to upsample.
-        low_res=((samples + 1) * 127.5).round() / 127.5 - 1,
+    # Create the classifier-free guidance tokens (empty)
+    full_batch_size = batch_size * 2
+    uncond_tokens, uncond_mask = model.tokenizer.padded_tokens_and_mask(
+        [], options['text_ctx']
+    )
 
-        # Text tokens
+    # Pack the tokens together into model kwargs.
+    model_kwargs = dict(
         tokens=th.tensor(
-            [tokens] * batch_size, device=device
+            [tokens] * batch_size + [uncond_tokens] * batch_size, device=device
         ),
         mask=th.tensor(
-            [mask] * batch_size,
+            [mask] * batch_size + [uncond_mask] * batch_size,
             dtype=th.bool,
             device=device,
         ),
     )
 
+    # Create a classifier-free guidance sampling function
+    def model_fn(x_t, ts, **kwargs):
+        half = x_t[: len(x_t) // 2]
+        combined = th.cat([half, half], dim=0)
+        model_out = model(combined, ts, **kwargs)
+        eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = th.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
+        eps = th.cat([half_eps, half_eps], dim=0)
+        return th.cat([eps, rest], dim=1)
+
     # Sample from the base model.
-    model_up.del_cache()
-    up_shape = (batch_size, 3, options_up["image_size"], options_up["image_size"])
-    up_samples = diffusion_up.ddim_sample_loop(
-        model_up,
-        up_shape,
-        noise=th.randn(up_shape, device=device) * upsample_temp,
+    model.del_cache()
+    samples = diffusion.p_sample_loop(
+        model_fn,
+        (full_batch_size, 3, options["image_size"], options["image_size"]),
         device=device,
         clip_denoised=True,
         progress=True,
         model_kwargs=model_kwargs,
         cond_fn=None,
     )[:batch_size]
-    model_up.del_cache()
+    model.del_cache()
 
     # Show the output
-    show_images(up_samples)
-demo = gr.Interface(fn =generate_upsampled_image_from_text,inputs ="text",outputs ="image")
-demo.launch()
+    show_images(samples)
+demo = gr.Interface(fn =generate_image_from_text,inputs ="text",outputs ="image")
+demo.launch()