Upload 2 files

app.py

@@ -68,21 +68,22 @@ 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)
+def gradio_generate(fibers_map: Image.Image,
+                    rings_map: Image.Image,
+                    num_steps: int) -> Image.Image:
+    # 1) uniformiser le mode
     fibers_map = fibers_map.convert("RGB")
     rings_map = rings_map.convert("RGB")
-    
-    result_img = inference(pipe, rings_map, fibers_map, num_steps)
-    
+
+
+    # 3) appeler l'inference avec la seed
+    result_img = inference(pipe,
+                           rings_map,
+                           fibers_map,
+                           num_steps)
     return result_img
 
-# --- 4) DÉFINITION DE L’INTERFACE GRADIO ---
+
 iface = gr.Interface(
     fn=gradio_generate,
     inputs=[
@@ -90,13 +91,17 @@ iface = gr.Interface(
         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"),
+    outputs=gr.Image(
+        type="pil",
+        label="Photorealistic wood generated",
+        format="png"           # ← force le .png au téléchargement
+    ),
     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.
 
@@ -104,10 +109,7 @@ iface = gr.Interface(
     """
 )
 
-# --- 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

@@ -12,6 +12,7 @@ import numpy as np
 import torch.nn as nn
 from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
 from PIL import Image
+import random
 
 class UNetNoCondWrapper(nn.Module):
     def __init__(self, base_unet: UNet2DModel):
@@ -41,8 +42,11 @@ class UNetNoCondWrapper(nn.Module):
         return self.unet.save_pretrained(save_directory, **kwargs)
     
 def inference(pipe, img1, img2, num_steps):
+    
+    seed = random.randrange(0, 2**32)
+    torch.manual_seed(seed)
         
-    generator = torch.Generator("cpu").manual_seed(0)   
+    generator = torch.Generator("cpu").manual_seed(seed)   
     
     img1 = img1.resize((512, 512))
     img2 = img2.resize((512, 512))
@@ -65,20 +69,6 @@ def inference(pipe, img1, img2, num_steps):
     
     all_images = []
     
-    # def cb_fn(step, timestep, latents):
-    #     # 1) Décoder
-    #     with torch.no_grad():
-    #         decoded_output = pipe.vae.decode(latents / pipe.vae.config.scaling_factor)
-    #     decoded_tensor = decoded_output.sample  # (B, C, H, W)
-        
-    #     # 2) Transformer en NumPy (channels last) et en uint8 [0–255]
-    #     t = decoded_tensor.cpu().clamp(0, 1)[0]           # (C, H, W)
-    #     arr = (t.permute(1, 2, 0).numpy() * 255).astype(np.uint8)  # (H, W, C)
-        
-    #     # 3) Créer la PIL.Image
-    #     img = Image.fromarray(arr)
-    #     all_images.append(img)
-    
     num_inference_steps = num_steps
     image_guidance_scale = 1.9
     guidance_scale = 10