import gradio as gr
import torch
import requests
from PIL import Image
import numpy as np
from torchvision.utils import make_grid, save_image
from diffusers import DiffusionPipeline

def load_wonder3d_pipeline():
    pipeline = DiffusionPipeline.from_pretrained(
        'flamehaze1115/wonder3d-v1.0',
        custom_pipeline='flamehaze1115/wonder3d-pipeline',
        torch_dtype=torch.float16
    )
    pipeline.unet.enable_xformers_memory_efficient_attention()
    if torch.cuda.is_available():
        pipeline.to('cuda:0')
    return pipeline

pipeline = load_wonder3d_pipeline()

def generate_3d(image):
    cond = Image.fromarray(np.array(image)[:, :, :3])
    images = pipeline(cond, num_inference_steps=20, output_type='pt', guidance_scale=1.0).images
    result = make_grid(images, nrow=6, ncol=2, padding=0, value_range=(0, 1))
    save_image(result, 'result.png')
    return 'result.png'

iface = gr.Interface(
    fn=generate_3d,
    inputs=gr.Image(),
    outputs=gr.Image(),
    title="Wonder3D Image to 3D",
    description="Загрузите изображение, чтобы сгенерировать 3D-модель."
)
iface.launch()