Update app.py

app.py

@@ -12,7 +12,9 @@ def infer(prompt):
     gen = torch.Generator(device=device)
 
     gen.manual_seed(21)
-    out = pipe(prompt=prompt, generator=gen, num_images_per_prompt=1, guidance_scale=7,
+    out = pipe(prompt=prompt, generator=gen, num_images_per_prompt=1, guidance_scale=7)
+    images = out.images[0]
+    out_edit = pipe(prompt=prompt, generator=gen, num_images_per_prompt=1, guidance_scale=7,
            editing_prompt=['male person',       # Concepts to apply 
                            'female person'],
            reverse_editing_direction=[True, False], # Direction of guidance i.e. decrease the first and increase the second concept
@@ -23,9 +25,9 @@ def infer(prompt):
            edit_mom_beta=0.6, # Momentum beta
            edit_weights=[1,1] # Weights of the individual concepts against each other
           )
-    image = out.images[0]
+    images_edited = out_edit.images[0]
     
-    return [image]
+    return [images, images_edited]
     
 
 css = """