Upload 2 files

app.py

@@ -68,46 +68,85 @@ pipe = StableDiffusionInstructPix2PixPipeline(
 
 pipe = pipe.to(torch.float32).to(device)
 
-
-
-# --- 3) FONCTION GRADIO D’INTERFACE ---
-def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image, num_steps): # -> 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")
+with gr.Blocks() as demo:
+    gr.Markdown("## Photorealistic Wood Generator\nUpload your two maps, run inference, then use the slider to browse steps.")
     
-    result_img = inference(pipe, device, rings_map, fibers_map, num_steps)
+    with gr.Row():
+        fibers = gr.Image(type="pil", label="Fibre orientation map")
+        rings  = gr.Image(type="pil", label="Growth ring map")
+    steps = gr.Number(value=10, label="Number of inference steps")
+    btn   = gr.Button("Generate")
+
+    # State pour stocker la liste des images
+    state_images = gr.State([])
+
+    # Slider pour parcourir
+    slider = gr.Slider(minimum=0, maximum=0, step=1, value=0, interactive=True, label="Step index")
+    # Image affichée
+    display = gr.Image(label="Intermediate result")
+
+    # 1) Au clique, on génère et on met à jour state + slider + display
+    def run_and_store(fib, ring, num_steps):
+        imgs = inference(fib, ring, int(num_steps))
+        # On renvoie : la liste, la nouvelle valeur max du slider, et l’image 0
+        return imgs, gr.Slider.update(maximum=len(imgs)-1, value=0), imgs[0]
+
+    btn.click(
+        fn=run_and_store,
+        inputs=[fibers, rings, steps],
+        outputs=[state_images, slider, display]
+    )
+
+    # 2) Quand on bouge le slider, on affiche state_images[slider]
+    def select_step(imgs, idx):
+        return imgs[int(idx)]
+
+    slider.change(
+        fn=select_step,
+        inputs=[state_images, slider],
+        outputs=display
+    )
+
+    demo.launch()
+
+# # --- 3) FONCTION GRADIO D’INTERFACE ---
+# def gradio_generate(fibers_map: Image.Image, rings_map: Image.Image, num_steps): # -> 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")
     
-    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"),
-        gr.Number(value=20, label="Number of inference steps")
-    ],
-    outputs=gr.Image(type="pil", label="Photorealistic wood generated"),
-    title="Photorealistic wood generator",
-    description="""
-    Upload :
-    1) a fibre orientation map,
-    2) a growth ring map.
+#     result_img = inference(pipe, rings_map, fibers_map, num_steps)
     
-    Set the number of inference steps.
-    Higher values can improve quality but increase processing time.
-
-    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)
+#     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"),
+#         gr.Number(value=10, label="Number of inference steps")
+#     ],
+#     outputs=gr.Image(type="pil", label="Photorealistic wood generated"),
+#     title="Photorealistic wood generator",
+#     description="""
+#     Upload :
+#     1) a fibre orientation map,
+#     2) a growth ring map.
+    
+#     Set the number of inference steps.
+#     Higher values can improve quality but increase processing time.
+
+#     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)

inference.py

@@ -40,7 +40,7 @@ class UNetNoCondWrapper(nn.Module):
         # délègue à la vraie instance UNet2DModel
         return self.unet.save_pretrained(save_directory, **kwargs)
     
-def inference(pipe,device, img1, img2, num_steps):
+def inference(pipe, img1, img2, num_steps):
         
     generator = torch.Generator("cpu").manual_seed(0)   
     
@@ -57,14 +57,22 @@ def inference(pipe,device, img1, img2, num_steps):
 
     img1_np[img1_np > 200] = 255
     img1_np[img1_np <= 200] = 0
-    img1_np = 255-img1_np
+    # 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)
+    
+    all_images = []
+    def cb_fn(step, timestep, latents):
+        # latents(torch.Tensor) -> image via VAE decode
+        with torch.no_grad():
+            decoded = pipe.vae.decode(latents / pipe.vae.config.scaling_factor)
+        # post‐traitement en PIL
+        img = pipe.numpy_to_pil(decoded.cpu().clamp(0,1))[0]
+        all_images.append(img)
 
     num_inference_steps = num_steps
-    print(num_inference_steps)
     image_guidance_scale = 1.9
     guidance_scale = 10
 
@@ -76,77 +84,14 @@ def inference(pipe,device, img1, img2, num_steps):
        guidance_scale=guidance_scale,
        generator=generator,
        safety_checker=None,
+       callback=cb_fn,
+       callback_steps=1,
        num_images_per_prompt=1
     ).images
     
-    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)
+    return all_images
+    # edited_image = edited_image[0]
     
-#     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]
+    # return edited_image
 
-#     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