Update app.py

app.py

@@ -35,7 +35,7 @@ guidance_scale = 3.0
 # A value of 1.0 is sharper, but sometimes results in grainy artifacts.
 upsample_temp = 0.997
 import gradio as gr
-def generate_image_from_text(prompt):
+def generate_upsampled_image_from_text(prompt):
     # Set the prompt text
     prompt = prompt
 
@@ -44,54 +44,43 @@ def generate_image_from_text(prompt):
     ##############################
 
     # Create the text tokens to feed to the model.
-    tokens = model.tokenizer.encode(prompt)
-    tokens, mask = model.tokenizer.padded_tokens_and_mask(
-        tokens, options['text_ctx']
+    tokens = model_up.tokenizer.encode(prompt)
+    tokens, mask = model_up.tokenizer.padded_tokens_and_mask(a
+        tokens, options_up['text_ctx']
     )
 
-    # 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']
-    )
-
-    # Pack the tokens together into model kwargs.
+    # Create the model conditioning dict.
     model_kwargs = dict(
+        # Low-res image to upsample.
+        low_res=((samples + 1) * 127.5).round() / 127.5 - 1,
+
+        # Text tokens
         tokens=th.tensor(
-            [tokens] * batch_size + [uncond_tokens] * batch_size, device=device
+            [tokens] * batch_size, device=device
         ),
         mask=th.tensor(
-            [mask] * batch_size + [uncond_mask] * batch_size,
+            [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.del_cache()
-    samples = diffusion.p_sample_loop(
-        model_fn,
-        (full_batch_size, 3, options["image_size"], options["image_size"]),
+    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,
         device=device,
         clip_denoised=True,
         progress=True,
         model_kwargs=model_kwargs,
         cond_fn=None,
     )[:batch_size]
-    model.del_cache()
+    model_up.del_cache()
 
     # Show the output
-    show_images(samples)
-demo = gr.Interface(fn =generate_image_from_text,inputs ="text",outputs ="image")
+    show_images(up_samples)
+demo = gr.Interface(fn =generate_upsampled_image_from_text,inputs ="text",outputs ="image")
 demo.launch()