Nikhil0987 commited on
Commit
efff157
1 Parent(s): 1418f33

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -11
app.py CHANGED
@@ -1,17 +1,35 @@
1
  # Text to 3D
 
 
2
  import torch
3
  from diffusers import ShapEPipeline
4
  from diffusers.utils import export_to_gif
5
 
6
- ckpt_id = "openai/shap-e" # Model identifier from the Hugging Face Model Hub
7
- pipe = ShapEPipeline.from_pretrained(ckpt_id).to("cuda") # Load the model correctly
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
 
9
- guidance_scale = 15.0
10
- prompt = "a shark"
11
- images = pipe(
12
- prompt,
13
- guidance_scale=guidance_scale,
14
- num_inference_steps=64,
15
- size=256,
16
- ).images
17
- gif_path = export_to_gif(images, "shark_3d.gif")
 
1
  # Text to 3D
2
+
3
+ import streamlit as st
4
  import torch
5
  from diffusers import ShapEPipeline
6
  from diffusers.utils import export_to_gif
7
 
8
+ # Model loading (Ideally done once at the start for efficiency)
9
+ ckpt_id = "openai/shap-e"
10
+ @st.cache_resource # Caches the model for faster subsequent runs
11
+ def load_model():
12
+ return ShapEPipeline.from_pretrained(ckpt_id).to("cuda")
13
+
14
+ pipe = load_model()
15
+
16
+ # App Title
17
+ st.title("Shark 3D Image Generator")
18
+
19
+ # User Inputs
20
+ prompt = st.text_input("Enter your prompt:", "a shark")
21
+ guidance_scale = st.slider("Guidance Scale", 0.0, 20.0, 15.0, step=0.5)
22
+
23
+ # Generate and Display Images
24
+ if st.button("Generate"):
25
+ with st.spinner("Generating images..."):
26
+ images = pipe(
27
+ prompt,
28
+ guidance_scale=guidance_scale,
29
+ num_inference_steps=64,
30
+ size=256,
31
+ ).images
32
+ gif_path = export_to_gif(images, "shark_3d.gif")
33
 
34
+ st.image(images[0]) # Display the first image
35
+ st.success("GIF saved as shark_3d.gif")