import gradio as gr def generate_from_latents(prompt, apply_edge_loss=False): # ... [All your processing, including the loop] # Generate image without edge loss image_without_loss = latents_to_pil(latents)[0] if apply_edge_loss: # Apply the edge (Laplacian) loss and regenerate the image # [Your logic for applying the edge loss and regenerating the image goes here] # You might run another loop similar to the one you have but with the edge loss applied. image_with_loss = latents_to_pil(latents_with_loss)[0] else: image_with_loss = image_without_loss # Or keep it the same if not applying edge loss return image_without_loss, image_with_loss # Define the Gradio interface iface = gr.Interface( fn=generate_from_latents, inputs=[ gr.inputs.Textbox(default="Canvas showing moonlight in night", label="Input Prompt"), gr.inputs.Checkbox(label="Apply Edge Loss") ], outputs=[ gr.outputs.Image(plot=True, label="Without Edge Loss"), gr.outputs.Image(plot=True, label="With Edge Loss") ] ) iface.launch()