Update app.py

app.py

@@ -51,44 +51,39 @@ def elapsed_time(fn, *args):
     return elapsed, output
 
 
-def main():
+st.title("BAYC Text to IMAGE generator")
+st.write(f"Stable diffusion model is fine-tuned by lora using dataset {DATASET_REPO}")
 
-    st.title("BAYC Text to IMAGE generator")
-    st.write(f"Stable diffusion model is fine-tuned by lora using dataset {DATASET_REPO}")
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+st.write("Loading models...")
+elapsed, pipeline = elapsed_time(load_pipeline_w_lora)
+st.write(f"Model is loaded in {elapsed} seconds!")
 
-    st.write("Loading models...")
-    elapsed, pipeline = elapsed_time(load_pipeline_w_lora)
-    st.write(f"Model is loaded in {elapsed} seconds!")
-    
-    pipeline = pipeline.to(device)
+pipeline = pipeline.to(device)
 
-    sample = Image.open(SAMPLE_IMAGE)
-    st.image(sample, caption="Example image with prompt <An ape with solid gold fur and beanie>")
+sample = Image.open(SAMPLE_IMAGE)
+st.image(sample, caption="Example image with prompt <An ape with solid gold fur and beanie>")
 
-    with st.form(key="information", clear_on_submit=True): 
-        prompt = st.text_input(
-            label="Write prompt to generate your unique BAYC image! (e.g. An ape with golden fur)")
-    
-        num_images = st.number_input(label="Number of images to generate", min_value=1, max_value=10)
-    
-        seed = st.number_input(label="Seed for images", min_value=1, max_value=10000)
+with st.form(key="information", clear_on_submit=True): 
+    prompt = st.text_input(
+        label="Write prompt to generate your unique BAYC image! (e.g. An ape with golden fur)")
 
-        submitted = st.form_submit_button(label="Submit")
-        
-    if submitted :
-        st.write(f"Generating {num_images} BAYC image with prompt <{prompt}>...")
+    num_images = st.number_input(label="Number of images to generate", min_value=1, max_value=10)
 
-        generator = torch.Generator(device=device).manual_seed(seed)
-        images = []
-        for img_idx in stqdm(range(num_images)):
-            generated_image = pipeline(prompt, num_inference_steps=30, generator=generator).images[0]
-            images.append(generated_image)
+    seed = st.number_input(label="Seed for images", min_value=1, max_value=10000)
+
+    submitted = st.form_submit_button(label="Submit")
+    
+if submitted :
+    st.write(f"Generating {num_images} BAYC image with prompt <{prompt}>...")
 
-        st.write("Done!")
+    generator = torch.Generator(device=device).manual_seed(seed)
+    images = []
+    for img_idx in stqdm(range(num_images)):
+        generated_image = pipeline(prompt, num_inference_steps=30, generator=generator).images[0]
+        images.append(generated_image)
 
-        st.image(images, width=150, caption=f"Generated Images with <{prompt}>")
+    st.write("Done!")
 
-if __name__ == '__main__':
-    main()
+    st.image(images, width=150, caption=f"Generated Images with <{prompt}>")