Spaces:

banyapon
/

Pana360gen

Build error

App Files Files Community

banyapon commited on Jul 27, 2024

Commit

e0be9d0

1 Parent(s): 647ac13

checkpoint 360 aframe fixed image

Browse files

Files changed (1) hide show

app.py +55 -23

app.py CHANGED Viewed

@@ -14,20 +14,17 @@ import spaces
 model_path = snapshot_download("archerfmy0831/sd-t2i-360panoimage")
 # Initialize pipelines
-pipe = Text2360PanoramaImagePipeline.from_pretrained(model_path, torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
-@spaces.GPU(duration=200)  # Assuming you have a GPU available
 def text_to_pano(prompt, upscale):
     try:
         input_data = {'prompt': prompt, 'upscale': upscale, 'refinement': False}
-        image = pipe(input_data).images[0]
-        return image
     except Exception as e:
         print(f"Error generating panorama: {e}")
-        return None
 title = """
 <div style="text-align: left;">
@@ -42,12 +39,12 @@ title = """
 </div>
 """
 def create_aframe_html(image):
     buffered = io.BytesIO()
     image.save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
-    return f"""
     <html>
     <head>
         <script src="https://aframe.io/releases/1.3.0/aframe.min.js"></script>
@@ -59,23 +56,58 @@ def create_aframe_html(image):
     </body>
     </html>
     """
 def create_aframe_preview(image):
     if image is None:
         return gr.HTML("Error generating image. Please check the prompt or try again.")
-    return gr.HTML(f"""
     <a-scene embedded style="width: 100%; height: 400px;">
-        <a-sky src="{image}" rotation="0 -130 0"></a-sky>
         <a-text font="kelsonsans" value="Create from AI" width="6" position="-2.5 0.25 -1.5" rotation="0 15 0"></a-text>
     </a-scene>
-    """)
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.HTML(title)
     with gr.Row():
         with gr.Column():
             t2p_input = gr.Textbox(label="Enter your prompt", lines=3)
@@ -83,27 +115,27 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             t2p_generate = gr.Button("Generate 360° Image")
         with gr.Column(variant="panel"):
             t2p_output = Pannellum(show_label=False, interactive=True)
     with gr.Row():
         with gr.Column():
-            t2p_image_output = gr.Image(label="Generated Image", type="pil")
             image_url_output = gr.Textbox(label="Image URL")
             gr.Button("Download Image").click(
                 None, inputs=[t2p_image_output], outputs=[t2p_image_output], _js="(img) => {window.open(img,'_blank');}"
             )
         with gr.Column(variant="panel"):
-            output_html = gr.Textbox(label="A-Frame HTML Code (Copy and use in Glitch)", lines=15)
-            gr.Markdown("Copy the HTML code below and paste it into a new Glitch project to view your 360° panorama in VR.")
-            with gr.Column():
-                aframe_preview = gr.HTML(label="A-Frame Preview")
     update_trigger = gr.State(value=0)
     def generate_with_update(prompt, upscale, trigger):
-        image = text_to_pano(prompt, upscale)
-        image_url = get_image_url(image)
         html = create_aframe_html(image) if image is not None else "Failed to generate A-Frame HTML. Please try again."
         preview = create_aframe_preview(image)
         return output, image, html, preview, image_url, trigger + 1
@@ -117,4 +149,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     image_url_output.change(fn=get_aframe_code, inputs=[image_url_output], outputs=[output_html])
-demo.launch()

 model_path = snapshot_download("archerfmy0831/sd-t2i-360panoimage")
 # Initialize pipelines
+txt2panoimg = Text2360PanoramaImagePipeline(model_path, torch_dtype=torch.float16)
+@spaces.GPU(duration=200)
 def text_to_pano(prompt, upscale):
     try:
         input_data = {'prompt': prompt, 'upscale': upscale, 'refinement': False}
+        output = txt2panoimg(input_data)
+        return output, output
     except Exception as e:
         print(f"Error generating panorama: {e}")
+        return None, None
 title = """
 <div style="text-align: left;">
 </div>
 """
 def create_aframe_html(image):
     buffered = io.BytesIO()
     image.save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    aframe_html = f"""
     <html>
     <head>
         <script src="https://aframe.io/releases/1.3.0/aframe.min.js"></script>
     </body>
     </html>
     """
+    return aframe_html
 def create_aframe_preview(image):
     if image is None:
         return gr.HTML("Error generating image. Please check the prompt or try again.")
+    buffered = io.BytesIO()
+    image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    aframe_html = f"""
     <a-scene embedded style="width: 100%; height: 400px;">
+        <a-assets>
+            <img id="pano-img" src="data:image/jpeg;base64,{img_str}">
+        </a-assets>
+        <a-sky src="#pano-img" rotation="0 -130 0"></a-sky>
         <a-text font="kelsonsans" value="Create from AI" width="6" position="-2.5 0.25 -1.5" rotation="0 15 0"></a-text>
     </a-scene>
+    """
+    return gr.HTML(aframe_html)
+def get_image_url(image):
+    if image is None:
+        return None
+    buffered = io.BytesIO()
+    image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return f"data:image/jpeg;base64,{img_str}"
+def get_aframe_code(image_url):
+    if image_url is None:
+        return "Error generating A-Frame code. Please generate an image first."
+    aframe_html = f"""
+    <html>
+    <head>
+        <script src="https://aframe.io/releases/1.3.0/aframe.min.js"></script>
+    </head>
+    <body>
+        <a-scene>
+            <a-sky src="{image_url}" rotation="0 -130 0"></a-sky>
+            <a-text font="kelsonsans" value="Create from AI" width="6" position="-2.5 0.25 -1.5" rotation="0 15 0"></a-text>
+        </a-scene>
+    </body>
+    </html>
+    """
+    return aframe_html
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.HTML(title)
     with gr.Row():
         with gr.Column():
             t2p_input = gr.Textbox(label="Enter your prompt", lines=3)
             t2p_generate = gr.Button("Generate 360° Image")
         with gr.Column(variant="panel"):
             t2p_output = Pannellum(show_label=False, interactive=True)
     with gr.Row():
         with gr.Column():
+            t2p_image_output = gr.Image(label="Generated Image")
             image_url_output = gr.Textbox(label="Image URL")
             gr.Button("Download Image").click(
                 None, inputs=[t2p_image_output], outputs=[t2p_image_output], _js="(img) => {window.open(img,'_blank');}"
             )
         with gr.Column(variant="panel"):
+            with gr.Column():
+                output_html = gr.Textbox(label="A-Frame HTML Code (Copy and use in Glitch)", lines=15)
+                gr.Markdown("Copy the HTML code below and paste it into a new Glitch project to view your 360° panorama in VR.")
+            with gr.Column(variant="panel"): # Add wrapper for A-Frame to avoid split screen
+                aframe_preview = gr.HTML(label="A-Frame Preview")
     update_trigger = gr.State(value=0)
     def generate_with_update(prompt, upscale, trigger):
+        output, image = text_to_pano(prompt, upscale)
+        image_url = get_image_url(image)
         html = create_aframe_html(image) if image is not None else "Failed to generate A-Frame HTML. Please try again."
         preview = create_aframe_preview(image)
         return output, image, html, preview, image_url, trigger + 1
     image_url_output.change(fn=get_aframe_code, inputs=[image_url_output], outputs=[output_html])
+demo.launch()