File size: 1,025 Bytes
0bc5ace f6b18bd 0bc5ace |
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 |
import torch
from diffusers import DiffusionPipeline
from PIL import Image
from utils import create_image_grid # Changed import
class MultiViewDiffusion:
def __init__(self, device="cuda", model_id = "dylanebert/mvdream"):
self.device = device
self.pipeline = DiffusionPipeline.from_pretrained(
model_id,
custom_pipeline="dylanebert/multi-view-diffusion",
torch_dtype=torch.float16,
trust_remote_code=True,
).to(self.device)
def generate_views(self, prompt, num_views=4, guidance_scale=5, num_inference_steps=30, elevation=0):
images = self.pipeline(
prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
elevation=elevation,
).images
return images
if __name__ == "__main__":
mv_diff = MultiViewDiffusion()
images = mv_diff.generate_views("A futuristic city")
grid = create_image_grid(images)
grid.save("multi_view_output.png") |