Update app.py

app.py

@@ -1,16 +1,22 @@
+import torch
 from diffusers import DiffusionPipeline
 import gradio as gr
 
 generator = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")
-generator.to("cuda")  # If you have a GPU, uncomment this line
+# move to GPU if available
+if torch.cuda.is_available():
+    generator = generator.to("cuda")
 
 def generate(prompts):
   images = generator(list(prompts)).images
   return [images]
 
-gr.Interface(generate, 
+demo = gr.Interface(generate, 
              "textbox", 
              "image", 
              batch=True, 
              max_batch_size=4  # Set the batch size based on your CPU/GPU memory
-).queue().launch()
+).queue()
+
+if __name__ == "__main__":
+    demo.launch()