Upload 2 files

app.py

@@ -9,67 +9,12 @@ import gradio as gr
 from PIL import Image
 import torch
 import torchvision.transforms as T
+from inference import inference
 
-# --- 1) IMPORTER VOTRE CODE D'INFERENCE ---
-# Par exemple, si vous avez un fichier inference.py qui définit une fonction `infer_wood(fibers_map, rings_map)`
-# vous pouvez faire :
-#
-# from inference import infer_wood
-#
-# Et vous assurez que `infer_wood` prend en entrée deux objets PIL.Image
-# (cartographie fibres et cartographie cernes) et renvoie une PIL.Image résultat.
-#
-# Si vous n'avez pas encore ce fichier, créez une fonction de type :
 
-# def infer_wood(fibers_img: Image.Image, rings_img: Image.Image) -> Image.Image:
-#     """
-#     Exemple de squelette de fonction d'inférence.
-#     -> Remplacez tout ce qui est à l'intérieur par votre propre pipeline (prétraitement, appel du modèle, post-traitement).
-#     """
-#     # --- Pré-traitement (adapté à votre modèle) ---
-#     # Par exemple :
-#     preprocess = T.Compose([
-#         T.Resize((256, 256)),
-#         T.ToTensor(),
-#         # T.Normalize(mean=[...], std=[...])  # si votre modèle a été entraîné avec normalisation
-#     ])
-#     x1 = preprocess(fibers_img).unsqueeze(0).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-#     x2 = preprocess(rings_img).unsqueeze(0).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-    
-#     # --- Chargement / usage du modèle (ici, c'est un exemple générique) ---
-#     # Imaginons que vous aviez déjà chargé votre modèle quelque part globalement :
-#     # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-#     # model = VotreModeleWood().to(device)
-#     # model.load_state_dict(torch.load('chemin/vers/votre_modele.pth', map_location=device))
-#     # model.eval()
-#     #
-#     # Ici, on concatène les deux cartes pour former l’entrée (adaptez selon votre archi).
-
-#     input_tensor = torch.cat([x1, x2], dim=1)  # par exemple (1, C1+C2, H, W)
-
-#     with torch.no_grad():
-#         pred_tensor = model(input_tensor)  # suppose que `model` est déjà défini globalement et chargé
-
-#     # --- Post-traitement pour revenir à PIL.Image ---
-#     postprocess = T.ToPILImage()
-#     output_img = postprocess(pred_tensor.squeeze(0).cpu().clamp(0, 1))
-#     return output_img
-
-def infer_wood(fibers_img: Image.Image, rings_img: Image.Image):    
-    return rings_img
-
-
-# --- 2) CHARGEMENT GLOBAL DU MODÈLE (optionnel) ---
-# Vous pouvez charger votre modèle une seule fois, ici, en dehors de la fonction infer_wood,
-# afin que Gradio ne fasse pas recharger à chaque appel. Par exemple :
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Exemple :
-# from models.votre_modele import VotreModeleWood
-# model = VotreModeleWood().to(device)
-# model.load_state_dict(torch.load('models/chemin_du_modele.pth', map_location=device))
-# model.eval()
 
 
 # --- 3) FONCTION GRADIO D’INTERFACE ---
@@ -81,8 +26,11 @@ def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image) -> Image.Im
     # 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")
-    # Appel de votre code d'inférence
-    result_img = infer_wood(fibers_map, rings_map)
+    
+    model_id = "CarolineM5/InstructPix2Pix_WithoutPrompt/model_LR"
+    
+    result_img = inference(model_id, device, rings_map, fibers_map)
+    
     return result_img
 
 
@@ -90,17 +38,17 @@ def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image) -> Image.Im
 iface = gr.Interface(
     fn=gradio_generate,
     inputs=[
-        gr.Image(type="pil", label="Cartographie d’orientation des fibres"),
-        gr.Image(type="pil", label="Cartographie des limites de cernes")
+        gr.Image(type="pil", label="Fibre orientation map"),
+        gr.Image(type="pil", label="Growth ring map")
     ],
-    outputs=gr.Image(type="pil", label="Bois photoréalistique généré"),
-    title="Générateur de bois photoréalistique",
+    outputs=gr.Image(type="pil", label="Photorealistic wood generated"),
+    title="Photorealistic wood generator",
     description="""
-    Téléversez :
-    1) une image de cartographie d’orientation des fibres,  
-    2) une image de cartographie des limites de cernes.  
+    Upload :
+    1) a fibre orientation mapping image,
+    2) a tree-ring boundary mapping image.
 
-    Le modèle renverra un rendu photoréalistique de bois que vous pouvez télécharger.
+    The model will return a photo-realistic rendering of the wood that you can download.
     """
 )
 

inference.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Jun 11 09:51:38 2025
+
+@author: camaac
+"""
+
+import PIL
+import torch
+from diffusers import StableDiffusionInstructPix2PixPipeline, UNet2DModel, AutoencoderKL, DDPMScheduler 
+import numpy as np
+import torch.nn as nn
+from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
+
+
+class UNetNoCondWrapper(nn.Module):
+    def __init__(self, base_unet: UNet2DModel):
+        super().__init__()
+        self.unet = base_unet
+
+    def forward(
+        self,
+        sample,
+        timestep,
+        encoder_hidden_states=None,    
+        added_cond_kwargs=None,        
+        cross_attention_kwargs=None,   
+        return_dict=False,
+        **kwargs                       
+    ):
+        
+        return self.unet(sample, timestep, return_dict=return_dict, **kwargs)
+
+    def __getattr__(self, name):
+        if name in ("unet", "forward", "__getstate__", "__setstate__"):
+            return super().__getattr__(name)
+        return getattr(self.unet, name)
+    
+    def save_pretrained(self, save_directory, **kwargs):
+        # délègue à la vraie instance UNet2DModel
+        return self.unet.save_pretrained(save_directory, **kwargs)
+    
+def inference(model_id,device, img1, img2):
+    
+
+    vae = AutoencoderKL.from_pretrained(f"{model_id}/vae").to(device)
+    scheduler = DDPMScheduler.from_pretrained(f"{model_id}/scheduler")
+    tokenizer = CLIPTokenizer.from_pretrained(f"{model_id}/tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(f"{model_id}/text_encoder").to(device)
+    feature_extractor = CLIPImageProcessor.from_pretrained(f"{model_id}/feature_extractor")
+    
+    # 2) Chargez votre UNet non‑conditionné et wrappez‑le
+    base_unet = UNet2DModel.from_pretrained(f"{model_id}/unet").to(device)
+    wrapped_unet = UNetNoCondWrapper(base_unet).to(device)
+    
+    # 3) Construisez la pipeline manuellement
+    pipe = StableDiffusionInstructPix2PixPipeline(
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=wrapped_unet,
+        scheduler=scheduler,
+        safety_checker=None,
+        feature_extractor=feature_extractor,
+    )
+    pipe = pipe.to(torch.float16).to(device)
+    
+    generator = torch.Generator("cuda").manual_seed(0)
+    
+    
+    img1 = img1.resize((512, 512))
+    img2 = img2.resize((512, 512))
+
+    img1_np = np.array(img1)
+    if len(img1_np.shape) > 2:
+        img1_np = img1_np[:, :, 0]
+            
+    img2_np = np.array(img2)
+    if len(img2_np.shape) > 2:
+        img2_np = img2_np[:, :, 0]
+
+    img1_np[img1_np > 200] = 255
+    img1_np[img1_np <= 200] = 0
+    img1_np = 255-img1_np
+    img_np = np.stack([img1_np, img2_np, img2_np], axis=2)
+
+    image = PIL.Image.fromarray(img_np)        
+    image = PIL.ImageOps.exif_transpose(image)
+
+    num_inference_steps = 20
+    image_guidance_scale = 1.9
+    guidance_scale = 10
+
+    edited_image = pipe(
+       prompt=[""] ,
+       image=image,
+       num_inference_steps=num_inference_steps,
+       image_guidance_scale=image_guidance_scale,
+       guidance_scale=guidance_scale,
+       generator=generator,
+       safety_checker=None,
+       num_images_per_prompt=1
+    ).images
+    
+    return edited_image