diffusers-workflow 0.4.0a3__py3-none-any.whl

dw/tasks/gather.py

@@ -0,0 +1,139 @@
+import glob as glob_lib
+import logging
+from diffusers.utils import load_image, load_video
+from ..arguments import fetch_image, fetch_video
+from ..security import validate_url, SecurityError
+
+logger = logging.getLogger("dw")
+
+
+def gather_images(glob=None, urls=None):
+    """
+    Gather images from local files and/or URLs.
+
+    Args:
+        glob: Pattern for matching local image files (e.g., "images/*.png")
+        urls: List of URLs to download images from
+
+    Returns:
+        List of loaded images
+
+    Raises:
+        ValueError: If no images are found
+        SecurityError: If validation fails
+    """
+    if urls is None:
+        urls = []
+    images = []
+
+    # Load local images matching glob pattern
+    if glob is not None:
+        logger.debug(f"Searching for images matching pattern: {glob}")
+        image_paths = glob_lib.glob(glob)
+        logger.info(f"Found {len(image_paths)} local images")
+
+        for path in image_paths:
+            try:
+                logger.debug(f"Loading image from: {path}")
+                images.append(fetch_image(path))
+            except SecurityError:
+                raise
+            except Exception as e:
+                logger.error(
+                    f"Failed to load image from {path}: {str(e)}", exc_info=True
+                )
+                raise
+
+    # Load images from URLs
+    for url in urls:
+        try:
+            logger.debug(f"Loading image from URL: {url}")
+            validated_url = validate_url(url)
+            images.append(load_image(validated_url))
+        except SecurityError:
+            raise
+        except Exception as e:
+            logger.error(
+                f"Failed to load image from URL {url}: {str(e)}", exc_info=True
+            )
+            raise
+
+    # Validate that we found at least one image
+    if len(images) == 0:
+        error_msg = "No images found"
+        logger.error(error_msg)
+        raise ValueError(error_msg)
+
+    logger.debug(f"Successfully gathered {len(images)} images")
+    return images
+
+
+def gather_videos(glob=None, urls=None):
+    """
+    Gather videos from local files and/or URLs.
+
+    Args:
+        glob: Pattern for matching local video files (e.g., "videos/*.mp4")
+        urls: List of URLs to download videos from
+
+    Returns:
+        List of loaded videos
+    """
+    if urls is None:
+        urls = []
+    videos = []
+
+    # Load local videos matching glob pattern
+    if glob is not None:
+        logger.debug(f"Searching for videos matching pattern: {glob}")
+        video_paths = glob_lib.glob(glob)
+        logger.info(f"Found {len(video_paths)} local videos")
+
+        for path in video_paths:
+            try:
+                logger.debug(f"Loading video from: {path}")
+                videos.append(fetch_video(path))
+            except SecurityError:
+                raise
+            except Exception as e:
+                logger.error(
+                    f"Failed to load video from {path}: {str(e)}", exc_info=True
+                )
+                raise
+
+    # Load videos from URLs
+    for url in urls:
+        try:
+            logger.debug(f"Loading video from URL: {url}")
+            validated_url = validate_url(url)
+            videos.append(load_video(validated_url))
+        except SecurityError:
+            raise
+        except Exception as e:
+            logger.error(
+                f"Failed to load video from URL {url}: {str(e)}", exc_info=True
+            )
+            raise
+
+    # Validate that we found at least one video
+    if len(videos) == 0:
+        error_msg = "No videos found"
+        logger.error(error_msg)
+        raise ValueError(error_msg)
+
+    logger.debug(f"Successfully gathered {len(videos)} videos")
+    return videos
+
+
+def gather_inputs(kwargs):
+    """
+    Gather input arguments for passing to next task.
+
+    Args:
+        kwargs: Dictionary of input arguments
+
+    Returns:
+        Input arguments unchanged
+    """
+    logger.debug(f"Gathering input arguments: {kwargs}")
+    return kwargs

dw/tasks/image_to_text.py

@@ -0,0 +1,43 @@
+"""
+Image captioning via HuggingFace transformers.
+
+Captioning is text generation that happens to be given a picture, so this
+delegates to the text_generation task rather than keeping a second copy of
+the pipeline handling. Transformers 5 removed the dedicated "image-to-text"
+task this used to build, and with it the BLIP-style captioning models; the
+work is done by vision-language models now.
+"""
+
+import logging
+from .text_generation import generate_text, _DEFAULT_VISION_MODEL
+
+logger = logging.getLogger("dw")
+
+_DEFAULT_MODEL = _DEFAULT_VISION_MODEL
+_DEFAULT_PROMPT = "Describe this image."
+
+
+def image_to_text(image, device="cpu", **kwargs):
+    """Generate a text caption for an image.
+
+    Args:
+        image: PIL Image (or URL/path) to caption.
+        device: Target device ("cuda", "mps", "cpu").
+        **kwargs:
+            model_name: HuggingFace model ID of a vision-language model
+                (default: HuggingFaceTB/SmolVLM-256M-Instruct).
+            prompt: What to ask about the image (default: "Describe this
+                image."). Ask a narrower question for a narrower caption.
+            system_prompt: Optional system instruction for the model.
+            max_new_tokens: Max tokens to generate (default: 50).
+
+    Returns:
+        Caption string.
+    """
+    prompt = kwargs.pop("prompt", _DEFAULT_PROMPT)
+    kwargs.setdefault("model_name", _DEFAULT_MODEL)
+    kwargs.setdefault("max_new_tokens", 50)
+
+    caption = generate_text(prompt, device=device, image=image, **kwargs)
+    logger.info(f"Caption: {caption[:100]}{'...' if len(caption) > 100 else ''}")
+    return caption