xinference 1.5.1__py3-none-any.whl → 1.6.0__py3-none-any.whl

xinference/core/media_interface.py

@@ -0,0 +1,758 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import io
+import logging
+import os
+import threading
+import time
+import uuid
+from typing import Dict, List, Optional, Tuple, Union
+
+import gradio as gr
+import PIL.Image
+from gradio import Markdown
+
+from ..client.restful.restful_client import (
+    RESTfulAudioModelHandle,
+    RESTfulImageModelHandle,
+    RESTfulVideoModelHandle,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MediaInterface:
+    def __init__(
+        self,
+        endpoint: str,
+        model_uid: str,
+        model_family: str,
+        model_name: str,
+        model_id: str,
+        model_revision: str,
+        model_ability: List[str],
+        model_type: str,
+        controlnet: Union[None, List[Dict[str, Union[str, None]]]],
+        access_token: Optional[str],
+    ):
+        self.endpoint = endpoint
+        self.model_uid = model_uid
+        self.model_family = model_family
+        self.model_name = model_name
+        self.model_id = model_id
+        self.model_revision = model_revision
+        self.model_ability = model_ability
+        self.model_type = model_type
+        self.controlnet = controlnet
+        self.access_token = (
+            access_token.replace("Bearer ", "") if access_token is not None else None
+        )
+
+    def build(self) -> gr.Blocks:
+        if self.model_type == "image":
+            assert "stable_diffusion" in self.model_family
+
+        interface = self.build_main_interface()
+        interface.queue()
+        # Gradio initiates the queue during a startup event, but since the app has already been
+        # started, that event will not run, so manually invoke the startup events.
+        # See: https://github.com/gradio-app/gradio/issues/5228
+        try:
+            interface.run_startup_events()
+        except AttributeError:
+            # compatibility
+            interface.startup_events()
+        favicon_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            os.path.pardir,
+            "web",
+            "ui",
+            "public",
+            "favicon.svg",
+        )
+        interface.favicon_path = favicon_path
+        return interface
+
+    def text2image_interface(self) -> "gr.Blocks":
+        from ..model.image.stable_diffusion.core import SAMPLING_METHODS
+
+        def text_generate_image(
+            prompt: str,
+            n: int,
+            size_width: int,
+            size_height: int,
+            guidance_scale: int,
+            num_inference_steps: int,
+            negative_prompt: Optional[str] = None,
+            sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
+        ) -> PIL.Image.Image:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert isinstance(model, RESTfulImageModelHandle)
+
+            size = f"{int(size_width)}*{int(size_height)}"
+            guidance_scale = None if guidance_scale == -1 else guidance_scale  # type: ignore
+            num_inference_steps = (
+                None if num_inference_steps == -1 else num_inference_steps  # type: ignore
+            )
+            sampler_name = None if sampler_name == "default" else sampler_name
+
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.text_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        n=n,
+                        size=size,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        negative_prompt=negative_prompt,
+                        sampler_name=sampler_name,
+                        response_format="b64_json",
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
+
+            images = []
+            for image_dict in response["data"]:  # type: ignore
+                assert image_dict["b64_json"] is not None
+                image_data = base64.b64decode(image_dict["b64_json"])
+                image = PIL.Image.open(io.BytesIO(image_data))
+                images.append(image)
+
+            return images
+
+        with gr.Blocks() as text2image_vl_interface:
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column(scale=10):
+                        prompt = gr.Textbox(
+                            label="Prompt",
+                            show_label=True,
+                            placeholder="Enter prompt here...",
+                        )
+                        negative_prompt = gr.Textbox(
+                            label="Negative prompt",
+                            show_label=True,
+                            placeholder="Enter negative prompt here...",
+                        )
+                    with gr.Column(scale=1):
+                        generate_button = gr.Button("Generate")
+
+                with gr.Row():
+                    n = gr.Number(label="Number of Images", value=1)
+                    size_width = gr.Number(label="Width", value=1024)
+                    size_height = gr.Number(label="Height", value=1024)
+                with gr.Row():
+                    guidance_scale = gr.Number(label="Guidance scale", value=-1)
+                    num_inference_steps = gr.Number(
+                        label="Inference Step Number", value=-1
+                    )
+                    sampler_name = gr.Dropdown(
+                        choices=SAMPLING_METHODS,
+                        value="default",
+                        label="Sampling method",
+                    )
+
+                with gr.Column():
+                    image_output = gr.Gallery()
+
+            generate_button.click(
+                text_generate_image,
+                inputs=[
+                    prompt,
+                    n,
+                    size_width,
+                    size_height,
+                    guidance_scale,
+                    num_inference_steps,
+                    negative_prompt,
+                    sampler_name,
+                ],
+                outputs=image_output,
+            )
+
+        return text2image_vl_interface
+
+    def image2image_interface(self) -> "gr.Blocks":
+        from ..model.image.stable_diffusion.core import SAMPLING_METHODS
+
+        def image_generate_image(
+            prompt: str,
+            negative_prompt: str,
+            image: PIL.Image.Image,
+            n: int,
+            size_width: int,
+            size_height: int,
+            num_inference_steps: int,
+            padding_image_to_multiple: int,
+            sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
+        ) -> PIL.Image.Image:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert isinstance(model, RESTfulImageModelHandle)
+
+            if size_width > 0 and size_height > 0:
+                size = f"{int(size_width)}*{int(size_height)}"
+            else:
+                size = None
+            num_inference_steps = (
+                None if num_inference_steps == -1 else num_inference_steps  # type: ignore
+            )
+            padding_image_to_multiple = None if padding_image_to_multiple == -1 else padding_image_to_multiple  # type: ignore
+            sampler_name = None if sampler_name == "default" else sampler_name
+
+            bio = io.BytesIO()
+            image.save(bio, format="png")
+
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.image_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        n=n,
+                        image=bio.getvalue(),
+                        size=size,
+                        response_format="b64_json",
+                        num_inference_steps=num_inference_steps,
+                        padding_image_to_multiple=padding_image_to_multiple,
+                        sampler_name=sampler_name,
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
+
+            images = []
+            for image_dict in response["data"]:  # type: ignore
+                assert image_dict["b64_json"] is not None
+                image_data = base64.b64decode(image_dict["b64_json"])
+                image = PIL.Image.open(io.BytesIO(image_data))
+                images.append(image)
+
+            return images
+
+        with gr.Blocks() as image2image_inteface:
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column(scale=10):
+                        prompt = gr.Textbox(
+                            label="Prompt",
+                            show_label=True,
+                            placeholder="Enter prompt here...",
+                        )
+                        negative_prompt = gr.Textbox(
+                            label="Negative Prompt",
+                            show_label=True,
+                            placeholder="Enter negative prompt here...",
+                        )
+                    with gr.Column(scale=1):
+                        generate_button = gr.Button("Generate")
+
+                with gr.Row():
+                    n = gr.Number(label="Number of image", value=1)
+                    size_width = gr.Number(label="Width", value=-1)
+                    size_height = gr.Number(label="Height", value=-1)
+
+                with gr.Row():
+                    num_inference_steps = gr.Number(
+                        label="Inference Step Number", value=-1
+                    )
+                    padding_image_to_multiple = gr.Number(
+                        label="Padding image to multiple", value=-1
+                    )
+                    sampler_name = gr.Dropdown(
+                        choices=SAMPLING_METHODS,
+                        value="default",
+                        label="Sampling method",
+                    )
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        uploaded_image = gr.Image(type="pil", label="Upload Image")
+                    with gr.Column(scale=1):
+                        output_gallery = gr.Gallery()
+
+            generate_button.click(
+                image_generate_image,
+                inputs=[
+                    prompt,
+                    negative_prompt,
+                    uploaded_image,
+                    n,
+                    size_width,
+                    size_height,
+                    num_inference_steps,
+                    padding_image_to_multiple,
+                    sampler_name,
+                ],
+                outputs=output_gallery,
+            )
+        return image2image_inteface
+
+    def text2video_interface(self) -> "gr.Blocks":
+        def text_generate_video(
+            prompt: str,
+            negative_prompt: str,
+            num_frames: int,
+            fps: int,
+            num_inference_steps: int,
+            guidance_scale: float,
+            width: int,
+            height: int,
+            progress=gr.Progress(),
+        ) -> List[Tuple[str, str]]:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert isinstance(model, RESTfulVideoModelHandle)
+
+            request_id = str(uuid.uuid4())
+            response = None
+            exc = None
+
+            # Run generation in a separate thread to allow progress tracking
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.text_to_video(
+                        request_id=request_id,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_frames=num_frames,
+                        fps=fps,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        width=width,
+                        height=height,
+                        response_format="b64_json",
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+
+            # Update progress bar during generation
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except Exception:
+                    cur_progress = 0.0
+                progress(cur_progress, desc="Generating video")
+                time.sleep(1)
+
+            if exc:
+                raise exc
+
+            # Decode and return the generated video
+            videos = []
+            for video_dict in response["data"]:  # type: ignore
+                video_data = base64.b64decode(video_dict["b64_json"])
+                video_path = f"/tmp/{uuid.uuid4()}.mp4"
+                with open(video_path, "wb") as f:
+                    f.write(video_data)
+                videos.append((video_path, "Generated Video"))
+
+            return videos
+
+        # Gradio UI definition
+        with gr.Blocks() as text2video_ui:
+            # Prompt & Negative Prompt (stacked vertically)
+            prompt = gr.Textbox(label="Prompt", placeholder="Enter video prompt")
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt", placeholder="Enter negative prompt"
+            )
+
+            # Parameters (2-column layout)
+            with gr.Row():
+                with gr.Column():
+                    width = gr.Number(label="Width", value=512)
+                    num_frames = gr.Number(label="Frames", value=16)
+                    steps = gr.Number(label="Inference Steps", value=25)
+                with gr.Column():
+                    height = gr.Number(label="Height", value=512)
+                    fps = gr.Number(label="FPS", value=8)
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale", minimum=1, maximum=20, value=7.5
+                    )
+
+            # Generate button
+            generate = gr.Button("Generate")
+
+            # Output gallery
+            gallery = gr.Gallery(label="Generated Videos", columns=2)
+
+            # Button click logic
+            generate.click(
+                fn=text_generate_video,
+                inputs=[
+                    prompt,
+                    negative_prompt,
+                    num_frames,
+                    fps,
+                    steps,
+                    guidance_scale,
+                    width,
+                    height,
+                ],
+                outputs=gallery,
+            )
+
+        return text2video_ui
+
+    def image2video_interface(self) -> "gr.Blocks":
+        def image_generate_video(
+            image: "PIL.Image",
+            prompt: str,
+            negative_prompt: str,
+            num_frames: int,
+            fps: int,
+            num_inference_steps: int,
+            guidance_scale: float,
+            width: int,
+            height: int,
+            progress=gr.Progress(),
+        ) -> List[Tuple[str, str]]:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert isinstance(model, RESTfulVideoModelHandle)
+
+            request_id = str(uuid.uuid4())
+            response = None
+            exc = None
+
+            # Convert uploaded image to base64
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+
+            # Run generation in a separate thread
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.image_to_video(
+                        request_id=request_id,
+                        image=buffered.getvalue(),
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_frames=num_frames,
+                        fps=fps,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        width=width,
+                        height=height,
+                        response_format="b64_json",
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+
+            # Progress loop
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except Exception:
+                    cur_progress = 0.0
+                progress(cur_progress, desc="Generating video from image")
+                time.sleep(1)
+
+            if exc:
+                raise exc
+
+            # Decode and return video files
+            videos = []
+            for video_dict in response["data"]:  # type: ignore
+                video_data = base64.b64decode(video_dict["b64_json"])
+                video_path = f"/tmp/{uuid.uuid4()}.mp4"
+                with open(video_path, "wb") as f:
+                    f.write(video_data)
+                videos.append((video_path, "Generated Video"))
+
+            return videos
+
+        # Gradio UI
+        with gr.Blocks() as image2video_ui:
+            image = gr.Image(label="Input Image", type="pil")
+
+            prompt = gr.Textbox(label="Prompt", placeholder="Enter video prompt")
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt", placeholder="Enter negative prompt"
+            )
+
+            with gr.Row():
+                with gr.Column():
+                    width = gr.Number(label="Width", value=512)
+                    num_frames = gr.Number(label="Frames", value=16)
+                    steps = gr.Number(label="Inference Steps", value=25)
+                with gr.Column():
+                    height = gr.Number(label="Height", value=512)
+                    fps = gr.Number(label="FPS", value=8)
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale", minimum=1, maximum=20, value=7.5
+                    )
+
+            generate = gr.Button("Generate")
+            gallery = gr.Gallery(label="Generated Videos", columns=2)
+
+            generate.click(
+                fn=image_generate_video,
+                inputs=[
+                    image,
+                    prompt,
+                    negative_prompt,
+                    num_frames,
+                    fps,
+                    steps,
+                    guidance_scale,
+                    width,
+                    height,
+                ],
+                outputs=gallery,
+            )
+
+        return image2video_ui
+
+    def audio2text_interface(self) -> "gr.Blocks":
+        def transcribe_audio(
+            audio_path: str,
+            language: Optional[str],
+            prompt: Optional[str],
+            temperature: float,
+        ) -> str:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert isinstance(model, RESTfulAudioModelHandle)
+
+            with open(audio_path, "rb") as f:
+                audio_data = f.read()
+
+            response = model.transcriptions(
+                audio=audio_data,
+                language=language or None,
+                prompt=prompt or None,
+                temperature=temperature,
+                response_format="json",
+            )
+
+            return response.get("text", "No transcription result.")
+
+        with gr.Blocks() as audio2text_ui:
+            with gr.Row():
+                audio_input = gr.Audio(
+                    type="filepath",
+                    label="Upload or Record Audio",
+                    sources=["upload", "microphone"],  # ✅ support both
+                )
+            with gr.Row():
+                language = gr.Textbox(
+                    label="Language", placeholder="e.g. en or zh", value=""
+                )
+                prompt = gr.Textbox(
+                    label="Prompt (optional)",
+                    placeholder="Provide context or vocabulary",
+                )
+                temperature = gr.Slider(
+                    label="Temperature", minimum=0.0, maximum=1.0, value=0.0, step=0.1
+                )
+            transcribe_btn = gr.Button("Transcribe")
+            output_text = gr.Textbox(label="Transcription", lines=5)
+
+            transcribe_btn.click(
+                fn=transcribe_audio,
+                inputs=[audio_input, language, prompt, temperature],
+                outputs=output_text,
+            )
+
+        return audio2text_ui
+
+    def text2speech_interface(self) -> "gr.Blocks":
+        def tts_generate(
+            input_text: str,
+            voice: str,
+            speed: float,
+            prompt_speech_file,
+            prompt_text: Optional[str],
+        ) -> str:
+            from ..client import RESTfulClient
+
+            client = RESTfulClient(self.endpoint)
+            client._set_token(self.access_token)
+            model = client.get_model(self.model_uid)
+            assert hasattr(model, "speech")
+
+            prompt_speech_bytes = None
+            if prompt_speech_file is not None:
+                with open(prompt_speech_file, "rb") as f:
+                    prompt_speech_bytes = f.read()
+
+            response = model.speech(
+                input=input_text,
+                voice=voice,
+                speed=speed,
+                response_format="mp3",
+                prompt_speech=prompt_speech_bytes,
+                prompt_text=prompt_text,
+            )
+
+            # Write to a temp .mp3 file and return its path
+            audio_path = f"/tmp/{uuid.uuid4()}.mp3"
+            with open(audio_path, "wb") as f:
+                f.write(response)
+
+            return audio_path
+
+        # Gradio UI
+        with gr.Blocks() as tts_ui:
+            with gr.Row():
+                with gr.Column():
+                    input_text = gr.Textbox(
+                        label="Text", placeholder="Enter text to synthesize"
+                    )
+                    voice = gr.Textbox(
+                        label="Voice", placeholder="Optional voice ID", value=""
+                    )
+                    speed = gr.Slider(
+                        label="Speed", minimum=0.5, maximum=2.0, value=1.0, step=0.1
+                    )
+
+                    prompt_speech = gr.Audio(
+                        label="Prompt Speech (for cloning)", type="filepath"
+                    )
+                    prompt_text = gr.Textbox(
+                        label="Prompt Text (for cloning)",
+                        placeholder="Text of the prompt speech",
+                    )
+
+                    generate = gr.Button("Generate")
+
+                with gr.Column():
+                    audio_output = gr.Audio(label="Generated Audio", type="filepath")
+
+            generate.click(
+                fn=tts_generate,
+                inputs=[input_text, voice, speed, prompt_speech, prompt_text],
+                outputs=audio_output,
+            )
+
+        return tts_ui
+
+    def build_main_interface(self) -> "gr.Blocks":
+        if self.model_type == "image":
+            title = f"🎨 Xinference Stable Diffusion: {self.model_name} 🎨"
+        elif self.model_type == "video":
+            title = f"🎨 Xinference Video Generation: {self.model_name} 🎨"
+        else:
+            assert self.model_type == "audio"
+            title = f"🎨 Xinference Audio Model: {self.model_name} 🎨"
+        with gr.Blocks(
+            title=title,
+            css="""
+                    .center{
+                        display: flex;
+                        justify-content: center;
+                        align-items: center;
+                        padding: 0px;
+                        color: #9ea4b0 !important;
+                    }
+                    """,
+            analytics_enabled=False,
+        ) as app:
+            Markdown(
+                f"""
+                    <h1 class="center" style='text-align: center; margin-bottom: 1rem'>{title}</h1>
+                    """
+            )
+            Markdown(
+                f"""
+                    <div class="center">
+                    Model ID: {self.model_uid}
+                    </div>
+                    """
+            )
+            if "text2image" in self.model_ability:
+                with gr.Tab("Text to Image"):
+                    self.text2image_interface()
+            if "image2image" in self.model_ability:
+                with gr.Tab("Image to Image"):
+                    self.image2image_interface()
+            if "text2video" in self.model_ability:
+                with gr.Tab("Text to Video"):
+                    self.text2video_interface()
+            if "image2video" in self.model_ability:
+                with gr.Tab("Image to Video"):
+                    self.image2video_interface()
+            if "audio2text" in self.model_ability:
+                with gr.Tab("Audio to Text"):
+                    self.audio2text_interface()
+            if "text2audio" in self.model_ability:
+                with gr.Tab("Text to Audio"):
+                    self.text2speech_interface()
+        return app