Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import PIL.Image | |
| import torch | |
| from diffusers import UniDiffuserPipeline | |
| class Model: | |
| def __init__(self): | |
| self.device = torch.device( | |
| 'cuda:0' if torch.cuda.is_available() else 'cpu') | |
| if self.device.type == 'cuda': | |
| self.pipe = UniDiffuserPipeline.from_pretrained( | |
| 'thu-ml/unidiffuser-v1', torch_dtype=torch.float16) | |
| self.pipe.to(self.device) | |
| else: | |
| self.pipe = UniDiffuserPipeline.from_pretrained( | |
| 'thu-ml/unidiffuser-v1') | |
| def run( | |
| self, | |
| mode: str, | |
| prompt: str, | |
| image: PIL.Image.Image | None, | |
| seed: int = 0, | |
| num_steps: int = 20, | |
| guidance_scale: float = 8.0, | |
| ) -> tuple[PIL.Image.Image | None, str]: | |
| generator = torch.Generator(device=self.device).manual_seed(seed) | |
| if mode == 't2i': | |
| self.pipe.set_text_to_image_mode() | |
| sample = self.pipe(prompt=prompt, | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return sample.images[0], '' | |
| elif mode == 'i2t': | |
| self.pipe.set_image_to_text_mode() | |
| sample = self.pipe(image=image, | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return None, sample.text[0] | |
| elif mode == 'joint': | |
| self.pipe.set_joint_mode() | |
| sample = self.pipe(num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return sample.images[0], sample.text[0] | |
| elif mode == 'i': | |
| self.pipe.set_image_mode() | |
| sample = self.pipe(num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return sample.images[0], '' | |
| elif mode == 't': | |
| self.pipe.set_text_mode() | |
| sample = self.pipe(num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return None, sample.text[0] | |
| elif mode == 'i2t2i': | |
| self.pipe.set_image_to_text_mode() | |
| sample = self.pipe(image=image, | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| self.pipe.set_text_to_image_mode() | |
| sample = self.pipe(prompt=sample.text[0], | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return sample.images[0], '' | |
| elif mode == 't2i2t': | |
| self.pipe.set_text_to_image_mode() | |
| sample = self.pipe(prompt=prompt, | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| self.pipe.set_image_to_text_mode() | |
| sample = self.pipe(image=sample.images[0], | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator) | |
| return None, sample.text[0] | |
| else: | |
| raise ValueError | |