update

.gitignore

@@ -0,0 +1,133 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+###
+.vscode/
+*.pth

README.md

@@ -1,14 +1,3 @@
----
-title: Control Animation
-emoji: 🔥
-sdk: gradio
-sdk_version: 3.23.0
-app_file: app.py
-pipeline_tag: text-to-video
-tags:
-- jax-diffusers-event
----
-
 # Control Animation
 
-Our code uses [Text2Video-Zero](https://github.com/Picsart-AI-Research/Text2Video-Zero) and the [Diffusers](https://github.com/huggingface/diffusers) library as inspiration.
+Our code uses [Text2Video-Zero](https://github.com/Picsart-AI-Research/Text2Video-Zero) and the [Diffusers](https://github.com/huggingface/diffusers) library as inspiration.

requirements.txt

@@ -1,6 +1,13 @@
-git+https://github.com/huggingface/diffusers.git #diffusers==0.16.0.dev0
--f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-jax[tpu]==0.4.5
+jax[cuda11_pip]
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jaxlib
+flax
+git+https://github.com/huggingface/diffusers@main
+opencv-python
+torch
+#git+https://github.com/huggingface/diffusers.git #diffusers==0.16.0.dev0
+#-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+#jax[tpu]==0.4.5
 absl-py==1.4.0
 accelerate==0.16.0
 addict==2.4.0
@@ -39,7 +46,6 @@ fastapi==0.95.1
 ffmpy==0.3.0
 filelock==3.11.0
 flatbuffers==23.3.3
-flax==0.6.7
 fonttools==4.39.3
 frozenlist==1.3.3
 fsspec==2023.4.0
@@ -60,8 +66,6 @@ imageio-ffmpeg==0.4.2
 importlib-metadata==6.5.0
 importlib-resources==5.12.0
 invisible-watermark==0.1.5
-jax
-jaxlib==0.4.4
 Jinja2==3.1.2
 joblib==1.2.0
 jsonschema==4.17.3
@@ -89,7 +93,7 @@ onnx==1.13.1
 onnxruntime==1.14.1
 open-clip-torch==2.16.0
 opencv-contrib-python==4.7.0.72
-opencv-python==4.7.0.72
+# opencv-python==4.7.0.72
 opencv-python-headless==4.7.0.72
 opt-einsum==3.3.0
 optax==0.1.4

text_to_animation/model.py

@@ -6,6 +6,9 @@ import jax.numpy as jnp
 import tomesd
 import jax
 
+from PIL import Image
+from typing import List
+
 from flax.training.common_utils import shard
 from flax.jax_utils import replicate
 from flax import jax_utils
@@ -381,31 +384,12 @@ class ControlAnimationModel:
             result, fps, path=path, watermark=gradio_utils.logo_name_to_path(watermark)
         )
 
-    def generate_animation(
-        self,
-        prompt: str,
-        model_link: str = "dreamlike-art/dreamlike-photoreal-2.0",
-        is_safetensor: bool = False,
-        motion_field_strength_x: int = 12,
-        motion_field_strength_y: int = 12,
-        t0: int = 44,
-        t1: int = 47,
-        n_prompt: str = "",
-        chunk_size: int = 8,
-        video_length: int = 8,
-        merging_ratio: float = 0.0,
-        seed: int = 0,
-        resolution: int = 512,
-        fps: int = 2,
-        use_cf_attn: bool = True,
-        use_motion_field: bool = True,
-        smooth_bg: bool = False,
-        smooth_bg_strength: float = 0.4,
-        path: str = None,
-    ):
-        if is_safetensor and model_link[-len(".safetensors") :] == ".safetensors":
-            pipe = utils.load_safetensors_model(model_link)
-        return
+    @staticmethod
+    def to_pil_images(images: torch.Tensor) -> List[Image.Image]:
+        images = (images / 2 + 0.5).clamp(0, 1)
+        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+        images = np.round(images * 255).astype(np.uint8)
+        return [Image.fromarray(image) for image in images]
 
     def generate_initial_frames(
         self,
@@ -419,9 +403,8 @@ class ControlAnimationModel:
         # batch_size: int = 1,
         cfg_scale: float = 7.0,
         seed: int = 0,
-    ):
-        print(f">>> prompt: {prompt}, model_link: {model_link}")
-
+    ) -> List[Image.Image]:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
         pipe = StableDiffusionPipeline.from_pretrained(model_link)
 
         batch_size = 4
@@ -434,6 +417,34 @@ class ControlAnimationModel:
             width=width,
             height=height,
             guidance_scale=cfg_scale,
+            generator=generator,
         ).images
+        pil_images = self.to_pil_images(images)
+
+        return pil_images
 
-        return images
+    def generate_animation(
+        self,
+        prompt: str,
+        model_link: str = "dreamlike-art/dreamlike-photoreal-2.0",
+        is_safetensor: bool = False,
+        motion_field_strength_x: int = 12,
+        motion_field_strength_y: int = 12,
+        t0: int = 44,
+        t1: int = 47,
+        n_prompt: str = "",
+        chunk_size: int = 8,
+        video_length: int = 8,
+        merging_ratio: float = 0.0,
+        seed: int = 0,
+        resolution: int = 512,
+        fps: int = 2,
+        use_cf_attn: bool = True,
+        use_motion_field: bool = True,
+        smooth_bg: bool = False,
+        smooth_bg_strength: float = 0.4,
+        path: str = None,
+    ):
+        if is_safetensor and model_link[-len(".safetensors") :] == ".safetensors":
+            pipe = utils.load_safetensors_model(model_link)
+        return

webui/app_control_animation.py

@@ -106,9 +106,7 @@ def create_demo(model: ControlAnimationModel):
                 with gr.Column(scale=3):
                     initial_frames = gr.Gallery(
                         label="Initial Frames", show_label=False
-                    ).style(
-                        columns=[2], rows=[2], object_fit="scale-down", height="auto"
-                    )
+                    ).style(columns=4, object_fit="contain")
                     initial_frames.select(select_initial_frame)
                     select_frame_button = gr.Button(
                         value="Select Initial Frame", variant="secondary"