Upload 2 files

app.py

@@ -9,13 +9,66 @@ import gradio as gr
 from PIL import Image
 import torch
 from inference import inference
+from diffusers import StableDiffusionInstructPix2PixPipeline, UNet2DModel, AutoencoderKL, DDPMScheduler 
+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)
+    
 # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 device = torch.device('cpu')
 
+model_id = "CarolineM5/InstructPix2Pix_WithoutPrompt"
+vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device)
+scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
+tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
+text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
+feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor")
+
+# 2) Chargez votre UNet non‑conditionné et wrappez‑le
+base_unet = UNet2DModel.from_pretrained(model_id, subfolder="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.float32).to(device)
+
+
+
 # --- 3) FONCTION GRADIO D’INTERFACE ---
 def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image) -> Image.Image:
     """
@@ -26,13 +79,10 @@ def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image) -> Image.Im
     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)
+    result_img = inference(pipe, device, rings_map, fibers_map)
     
     return result_img
 
-
 # --- 4) DÉFINITION DE L’INTERFACE GRADIO ---
 iface = gr.Interface(
     fn=gradio_generate,
@@ -44,8 +94,8 @@ iface = gr.Interface(
     title="Photorealistic wood generator",
     description="""
     Upload :
-    1) a fibre orientation mapping image,
-    2) a tree-ring boundary mapping image.
+    1) a fibre orientation map,
+    2) a growth ring map.
 
     The model will return a photo-realistic rendering of the wood that you can download.
     """

inference.py

@@ -40,34 +40,9 @@ class UNetNoCondWrapper(nn.Module):
         # 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(model_id, subfolder="vae").to(device)
-    scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
-    tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
-    text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
-    feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor")
-    
-    # 2) Chargez votre UNet non‑conditionné et wrappez‑le
-    base_unet = UNet2DModel.from_pretrained(model_id, subfolder="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)
-    pipe = pipe.to(torch.float32).to(device)
-    
-    generator = torch.Generator("cpu").manual_seed(0)
-    
+def inference(pipe,device, img1, img2):
+        
+    generator = torch.Generator("cpu").manual_seed(0)   
     
     img1 = img1.resize((512, 512))
     img2 = img2.resize((512, 512))
@@ -88,7 +63,7 @@ def inference(model_id,device, img1, img2):
     image = PIL.Image.fromarray(img_np)        
     image = PIL.ImageOps.exif_transpose(image)
 
-    num_inference_steps = 20
+    num_inference_steps = 5
     image_guidance_scale = 1.9
     guidance_scale = 10
 
@@ -106,3 +81,71 @@ def inference(model_id,device, img1, img2):
     edited_image = edited_image[0]
     
     return edited_image
+
+
+# def inference(model_id,device, img1, img2):
+    
+    
+#     vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device)
+#     scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
+#     tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
+#     text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
+#     feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor")
+    
+#     # 2) Chargez votre UNet non‑conditionné et wrappez‑le
+#     base_unet = UNet2DModel.from_pretrained(model_id, subfolder="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)
+#     pipe = pipe.to(torch.float32).to(device)
+    
+#     generator = torch.Generator("cpu").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
+    
+#     edited_image = edited_image[0]
+    
+#     return edited_image