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