import streamlit as st
from PIL import Image
import torch
from diffusers import DiffusionPipeline  # تأكد من استخدام مكتبة متوافقة مع النموذج الجديد
import numpy as np
import io

# إعداد الصفحة
st.set_page_config(page_title="3D Mesh Generation", page_icon="🖼️")

# عنوان الصفحة
st.title("3D Mesh Generation from Image")

# تعليمات
st.write("Upload an image to generate a 3D mesh.")

# تحميل النموذج
@st.cache_resource
def load_model():
    model_id = "your_new_model_id_here"  # استبدل بـ ID النموذج الجديد
    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
    return pipe

pipe = load_model()

# إدخال المستخدم
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_image is not None:
    image = Image.open(uploaded_image).convert("RGB")
    st.image(image, caption='Uploaded Image', use_column_width=True)

    if st.button('Generate 3D Mesh'):
        # تحويل الصورة إلى تنسور
        image_tensor = torch.tensor(np.array(image)).float().unsqueeze(0).permute(0, 3, 1, 2).to("cuda") / 255.0
        
        # توليد الشبكة ثلاثية الأبعاد
        with st.spinner("Generating 3D mesh..."):
            mesh = pipe(image_tensor).images  # تأكد من أن هذا هو الاستخدام الصحيح للنموذج

        st.success("3D mesh generated successfully!")

        # عرض النتيجة
        st.write("### 3D Mesh Preview")
        st.image(mesh[0], caption="Generated 3D Mesh", use_column_width=True)

        # تنزيل النموذج
        mesh_bytes = io.BytesIO()
        mesh[0].save(mesh_bytes, format="png")  # استخدم الصيغة المناسبة هنا
        st.download_button(label="Download 3D Mesh", data=mesh_bytes.getvalue(), file_name="generated_mesh.png", mime="image/png")

# تقديم بعض المعلومات حول النموذج
st.write("""
### About the Model:
This model generates 3D meshes from single images