Update app.py

app.py

@@ -110,7 +110,6 @@
 #         st.image(pil_images[0]) 
 #         st.success("GIF saved as shark_3d.gif")
 
-
 import streamlit as st
 import torch
 from diffusers import StableDiffusionPipeline
@@ -120,11 +119,10 @@ import numpy as np
 
 # Model loading 
 ckpt_id = "openai/shap-e"
-StableDiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=torch.float16, use_auth_token=True).to("cuda") 
 @st.cache_resource
 def load_model():
-    return StableDiffusionPipeline.from_pretrained(ckpt_id).to("cuda")  
-pipe = load_model()
+    return StableDiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=torch.float16, use_auth_token=True).to("cuda")  
+pipe = load_model()  # Load the model 
 
 # App Title
 st.title("Shark 3D Image Generator")
@@ -140,16 +138,19 @@ def should_resize(image):
         return False 
 
 # Generate and Display Images
-if st.button("Generate"):    
-    with st.spinner("Generating images..."):
-        images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64).images
+if st.button("Generate"):        
+    with st.spinner("Generating images..."): 
+        images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64).images 
 
         pil_images = [] 
-        for image in images:            
+        for image in images:
+            # ***Insert image conversion logic here*** 
+
             if should_resize(image): 
                 image = image.resize((256, 256))                
             pil_images.append(image) 
 
         gif_path = export_to_gif(pil_images, "shark_3d.gif")         
-        st.image(pil_images[0])         
+        st.image(pil_images[0])          
         st.success("GIF saved as shark_3d.gif")
+