Update app.py

app.py

@@ -1,23 +1,27 @@
+import gradio as gr
+import spaces
 import torch
 from diffusers import DiffusionPipeline
-import huggingface_hub
-import requests
 from PIL import Image
-from io import BytesIO
-import numpy as np
-import gradio as gr
 
-multi_view_diffusion_pipeline = DiffusionPipeline.from_pretrained(
-    "reiirene/multi-view-diffusion",
-    custom_pipeline="reiirene/multi-view-diffusion",
+
+# Text-to-Multi-View Diffusion pipeline
+text_pipeline = DiffusionPipeline.from_pretrained(
+    "dylanebert/mvdream",
+    custom_pipeline="dylanebert/multi-view-diffusion",
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+).to("cuda")
+
+
+# Image-to-Multi-View Diffusion pipeline
+image_pipeline = DiffusionPipeline.from_pretrained(
+    "dylanebert/multi-view-diffusion",
+    custom_pipeline="dylanebert/multi-view-diffusion",
     torch_dtype=torch.float16,
     trust_remote_code=True,
 ).to("cuda")
 
-image_url = "https://huggingface.co/datasets/dylanebert/3d-arena/resolve/main/inputs/images/a_cat_statue.jpg"
-response = requests.get(image_url)
-image = Image.open(BytesIO(response.content))
-image
 
 def create_image_grid(images):
     images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
@@ -32,26 +36,53 @@ def create_image_grid(images):
 
     return grid_img
 
-image = np.array(image, dtype=np.float32) / 255.0
-images = multi_view_diffusion_pipeline("", image, guidance_scale=5, num_inference_steps=30, elevation=0)
 
-create_image_grid(images)
+@spaces.GPU
+def text_to_mv(prompt):
+    images = text_pipeline(
+        prompt, guidance_scale=5, num_inference_steps=30, elevation=0
+    )
+    return create_image_grid(images)
+
 
-def run(image):
-  image = np.array(image, dtype=np.float32) / 255.0
-  images = multi_view_diffusion_pipeline("", image, guidance_scale=5, num_inference_steps=30, elevation=0)
+@spaces.GPU
+def image_to_mv(image, prompt):
+    image = image.astype("float32") / 255.0
+    images = image_pipeline(
+        prompt, image, guidance_scale=5, num_inference_steps=30, elevation=0
+    )
+    return create_image_grid(images)
 
-  images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
 
-  width, height = images[0].size
-  grid_img = Image.new("RGB", (2 * width, 2 * height))
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            with gr.Tab("Text Input"):
+                text_input = gr.Textbox(
+                    lines=2,
+                    show_label=False,
+                    placeholder="Enter a prompt here (e.g. 'a cat statue')",
+                )
+                text_btn = gr.Button("Generate Multi-View Images")
+            with gr.Tab("Image Input"):
+                image_input = gr.Image(
+                    label="Image Input",
+                    type="numpy",
+                )
+                optional_text_input = gr.Textbox(
+                    lines=2,
+                    show_label=False,
+                    placeholder="Enter an optional prompt here",
+                )
+                image_btn = gr.Button("Generate Multi-View Images")
+        with gr.Column():
+            output = gr.Image(label="Generated Images")
 
-  grid_img.paste(images[0], (0, 0))
-  grid_img.paste(images[1], (width, 0))
-  grid_img.paste(images[2], (0, height))
-  grid_img.paste(images[3], (width, height))
+    text_btn.click(fn=text_to_mv, inputs=text_input, outputs=output)
+    image_btn.click(
+        fn=image_to_mv, inputs=[image_input, optional_text_input], outputs=output
+    )
 
-  return grid_img
 
-demo = gr.Interface(fn=run, inputs="image", outputs="image")
-demo.launch(debug=True)
+if __name__ == "__main__":
+    demo.queue().launch()