Create main.py

main.py

@@ -0,0 +1,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,
+                }
+            }
+        }
+    )