Update app.py

app.py

@@ -204,11 +204,56 @@ class ImStack(nn.Module):
  
 
 
-def generate(text, n_steps):
-  # Encode prompt
-  embed = perceptor.encode_text(clip.tokenize(text).to(device)).float()
-  #todo
-  return np.random.random((128, 128, 3)).astype(np.uint8)
+def generate(text, n_iter):
+    
+    lr=0.25 #@param
+    # init_image=None #@param
+    weight_decay=1e-5 #@param
+    out_size=540 #@param
+    base_size=20 #@param
+    n_layers=4 #@param
+    scale=3 #@param
+
+
+
+    p_prompts = []
+    embed = perceptor.encode_text(clip.tokenize(text).to(device)).float()
+    p_prompts.append(Prompt(embed, 1, float('-inf')).to(device)) # 1 is the weight
+
+    # SOme negative prompts
+    n_prompts = []
+    for pr in ["Random noise", 'saturated rainbow RGB deep dream']:
+      embed = perceptor.encode_text(clip.tokenize(pr).to(device)).float()
+      n_prompts.append(Prompt(embed, 0.5, float('-inf')).to(device)) # 0.5 is the weight
+
+    # The ImageStack - trying a different scale and n_layers
+    ims = ImStack(base_size=base_size, scale=scale, n_layers=n_layers, out_size=out_size, decay=0.4)
+    
+    optimizer = optim.Adam(ims.layers, lr=lr, weight_decay=weight_decay)
+    losses = []
+
+    for i in range(n_iter):
+        optimizer.zero_grad()
+    
+        if i < 15: # Save time by skipping the cutouts and focusing on the lower layers 
+          im = ims.preview(n_preview=1 + i//20 )
+          iii = perceptor.encode_image(normalize(im)).float()
+        else:
+          im = ims()
+          iii = perceptor.encode_image(normalize(make_cutouts(im))).float()
+        
+        l = 0
+        for prompt in p_prompts:
+          l += prompt(iii)
+        for prompt in n_prompts:
+          l -= prompt(iii)
+    
+        losses.append(float(l.detach().cpu()))
+        l.backward() # Backprop
+        optimizer.step() # Update
+
+  im = ims.to_pil() 
+  return np.array(im)
 
 iface = gr.Interface(fn=generate, 
   inputs=[