Upload 2 files

app.py

@@ -71,7 +71,7 @@ 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:
+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).
@@ -80,7 +80,7 @@ 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")
     
-    result_img = inference(pipe, device, rings_map, fibers_map)
+    result_img = inference(pipe, device, rings_map, fibers_map, num_steps)
     
     return result_img
 
@@ -89,7 +89,8 @@ iface = gr.Interface(
     fn=gradio_generate,
     inputs=[
         gr.Image(type="pil", label="Fibre orientation map"),
-        gr.Image(type="pil", label="Growth ring 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",
@@ -97,6 +98,9 @@ iface = gr.Interface(
     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.
     """

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):
+def inference(pipe,device, img1, img2, num_steps):
         
     generator = torch.Generator("cpu").manual_seed(0)   
     
@@ -63,7 +63,7 @@ def inference(pipe,device, img1, img2):
     image = PIL.Image.fromarray(img_np)        
     image = PIL.ImageOps.exif_transpose(image)
 
-    num_inference_steps = 5
+    num_inference_steps = num_steps
     image_guidance_scale = 1.9
     guidance_scale = 10