app.py

# -*- coding: utf-8 -*-
"""
Created on Tue Jun 10 11:16:28 2025

@author: camaac
"""

import gradio as gr
from PIL import Image
import torch
from inference import inference



device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')



# --- 3) FONCTION GRADIO D’INTERFACE ---
def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image) -> Image.Image:
    """
    Cette fonction sera appelée à chaque upload par Gradio.
    Elle doit retourner une PIL.Image (ou un chemin vers l’image enregistrée).
    """
    # Vérifier que les deux images sont bien en mode RGB (ou adapter si besoin)
    fibers_map = fibers_map.convert("RGB")
    rings_map = rings_map.convert("RGB")
    
    model_id = "CarolineM5/InstructPix2Pix_WithoutPrompt"
    
    result_img = inference(model_id, device, rings_map, fibers_map)
    
    return result_img


# --- 4) DÉFINITION DE L’INTERFACE GRADIO ---
iface = gr.Interface(
    fn=gradio_generate,
    inputs=[
        gr.Image(type="pil", label="Fibre orientation map"),
        gr.Image(type="pil", label="Growth ring map")
    ],
    outputs=gr.Image(type="pil", label="Photorealistic wood generated"),
    title="Photorealistic wood generator",
    description="""
    Upload :
    1) a fibre orientation mapping image,
    2) a tree-ring boundary mapping image.

    The model will return a photo-realistic rendering of the wood that you can download.
    """
)

# --- 5) LANCER L’APPLICATION ---
if __name__ == "__main__":
    # Vous pouvez préciser `server_name="0.0.0.0"` si vous souhaitez qu’il soit accessible sur le réseau
    # et `server_port=7860` (ou autre port) si vous voulez le personnaliser.
    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)