Update app_main.py

app_main.py

@@ -1,4 +1,4 @@
-from diffusers import StableDiffusionLDM3DPipeline
+from optimum.habana.diffusers import GaudiDDIMScheduler, GaudiStableDiffusionLDM3DPipeline
 import gradio as gr
 import torch 
 from PIL import Image
@@ -8,19 +8,16 @@ from tempfile import NamedTemporaryFile
 from pathlib import Path
 
 Path("tmp").mkdir(exist_ok=True)
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "hpu"
 print(f"Device is {device}")
-torch_type = torch.float16 if device == "cuda" else torch.float32
-pipe = StableDiffusionLDM3DPipeline.from_pretrained(
-    "Intel/ldm3d-pano", 
-    torch_dtype=torch_type
-    # , safety_checker=None
-)
+model_name = "Intel/ldm3d-pano"
+scheduler = GaudiDDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
+pipe = GaudiStableDiffusionLDM3DPipeline.from_pretrained( model_name,
+                        scheduler=scheduler,
+                        use_habana=True,
+                        use_hpu_graphs=True,
+                        gaudi_config="Habana/stable-diffusion")
 pipe.to(device)
-if device == "cuda":
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe.enable_model_cpu_offload()
-
 
 def get_iframe(rgb_path: str, depth_path: str, viewer_mode: str = "6DOF"):
     # buffered = BytesIO()