Spaces:
Runtime error
Runtime error
Create main.py
Browse files
main.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
from streamlit_pannellum import streamlit_pannellum
|
| 3 |
+
from diffusers import StableDiffusionLDM3DPipeline
|
| 4 |
+
from PIL import Image
|
| 5 |
+
from typing import Optional
|
| 6 |
+
from torch import Tensor
|
| 7 |
+
from torch.nn import functional as F
|
| 8 |
+
from torch.nn import Conv2d
|
| 9 |
+
from torch.nn.modules.utils import _pair
|
| 10 |
+
|
| 11 |
+
# Function to override _conv_forward method
|
| 12 |
+
def asymmetricConv2DConvForward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):
|
| 13 |
+
paddingX = (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0)
|
| 14 |
+
paddingY = (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3])
|
| 15 |
+
working = F.pad(input, paddingX, mode='circular')
|
| 16 |
+
working = F.pad(working, paddingY, mode='constant')
|
| 17 |
+
return F.conv2d(working, weight, bias, self.stride, _pair(0), self.dilation, self.groups)
|
| 18 |
+
|
| 19 |
+
# Load the pipeline
|
| 20 |
+
pipe = StableDiffusionLDM3DPipeline.from_pretrained("Intel/ldm3d-pano")
|
| 21 |
+
pipe.to("cuda")
|
| 22 |
+
|
| 23 |
+
# Patch the Conv2d layers
|
| 24 |
+
targets = [pipe.vae, pipe.text_encoder, pipe.unet]
|
| 25 |
+
for target in targets:
|
| 26 |
+
for module in target.modules():
|
| 27 |
+
if isinstance(module, Conv2d):
|
| 28 |
+
module._conv_forward = asymmetricConv2DConvForward.__get__(module, Conv2d)
|
| 29 |
+
|
| 30 |
+
# Function to generate panoramic images
|
| 31 |
+
def generate_panoramic_image(prompt, name):
|
| 32 |
+
output = pipe(prompt, width=1024, height=512, guidance_scale=7.0, num_inference_steps=50)
|
| 33 |
+
rgb_image, depth_image = output.rgb, output.depth
|
| 34 |
+
rgb_image[0].save(name + "_ldm3d_rgb.jpg")
|
| 35 |
+
depth_image[0].save(name + "_ldd3d_depth.png")
|
| 36 |
+
return name + "_ldm3d_rgb.jpg", name + "_ldd3d_depth.png"
|
| 37 |
+
|
| 38 |
+
# Streamlit Interface
|
| 39 |
+
st.title("Pannellum Streamlit plugin")
|
| 40 |
+
st.markdown("This space is a showcase of the [streamlit_pannellum](https://gitlab.com/nicolalandro/streamlit-pannellum) lib.")
|
| 41 |
+
|
| 42 |
+
prompt = st.text_input("Enter a prompt for the panoramic image",
|
| 43 |
+
"360, Ben Erdt, Ognjen Sporin, Raphael Lacoste. A garden of oversized flowers...")
|
| 44 |
+
|
| 45 |
+
generate_button = st.button("Generate Panoramic Image")
|
| 46 |
+
|
| 47 |
+
if generate_button:
|
| 48 |
+
name = "generated_image" # This can be dynamic
|
| 49 |
+
rgb_image_path, _ = generate_panoramic_image(prompt, name)
|
| 50 |
+
|
| 51 |
+
# Display the generated panoramic image in Pannellum viewer
|
| 52 |
+
streamlit_pannellum(
|
| 53 |
+
config={
|
| 54 |
+
"default": {
|
| 55 |
+
"firstScene": "generated",
|
| 56 |
+
"autoLoad": True
|
| 57 |
+
},
|
| 58 |
+
"scenes": {
|
| 59 |
+
"generated": {
|
| 60 |
+
"title": "Generated Panoramic Image",
|
| 61 |
+
"type": "equirectangular",
|
| 62 |
+
"panorama": rgb_image_path,
|
| 63 |
+
"autoLoad": True,
|
| 64 |
+
}
|
| 65 |
+
}
|
| 66 |
+
}
|
| 67 |
+
)
|