Spaces:
Sleeping
Sleeping
| # -*- 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 | |
| from PIL import Image | |
| import random | |
| 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(pipe, img1, img2, num_steps): | |
| seed = random.randrange(0, 2**32) | |
| torch.manual_seed(seed) | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| 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) | |
| all_images = [] | |
| num_inference_steps = num_steps | |
| 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].convert("L") | |
| return edited_image | |