Update app.py

app.py

@@ -43,40 +43,13 @@ 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"
 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
-##############################
-# Sample from the base model #
-##############################
 
-# 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']
-)
-
-# 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.
-model_kwargs = dict(
-    tokens=th.tensor(
-        [tokens] * batch_size + [uncond_tokens] * batch_size, device=device
-    ),
-    mask=th.tensor(
-        [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):
@@ -89,7 +62,35 @@ def model_fn(x_t, ts, **kwargs):
     eps = th.cat([half_eps, half_eps], dim=0)
     return th.cat([eps, rest], dim=1)
     
-def run():    
+def run(prompt):    
+
+    ##############################
+    # Sample from the base model #
+    ##############################
+    
+    # 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']
+    )
+    
+    # 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.
+    model_kwargs = dict(
+        tokens=th.tensor(
+            [tokens] * batch_size + [uncond_tokens] * batch_size, device=device
+        ),
+        mask=th.tensor(
+            [mask] * batch_size + [uncond_mask] * batch_size,
+            dtype=th.bool,
+            device=device,
+        ),
+)
 
 
     print('run():')