Flyte

"""Run high-throughout inference for generative AI."""
  
import base64
import torch

import flyte
import flyte.io
import flyte.report

env = flyte.TaskEnvironment(
    name="stable-diffusion",
    image=(
        flyte.Image.from_debian_base()
        .with_commands(["pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124"])
        .with_requirements("requirements.txt")
    ),
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu=1),
)

@env.task(report=True)
async def generate(prompt: str, steps: int = 30) -> flyte.io.File:
    """Generate an image from a text prompt using Stable Diffusion."""
    from diffusers import AutoPipelineForText2Image

    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        variant="fp16" if device == "cuda" else None,
    ).to(device)

    image = pipe(prompt, num_inference_steps=steps, guidance_scale=0.0).images[0]

    path = "/tmp/output.png"
    image.save(path)

    with open(path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()
    await flyte.report.replace.aio(
        f"<h2>Stable Diffusion</h2>"
        f"<p><b>Prompt:</b> {prompt}</p>"
        f'<img src="data:image/png;base64,{img_b64}" style="max-width:512px" />'
    )
    await flyte.report.flush.aio()

    return await flyte.io.File.from_local(path)

#flyte run stable_diffusion.py generate --prompt "a cat astronaut floating in space, digital art"

"""Run data ETL jobs at scale with modern frameworks."""
  
import duckdb
import pandas as pd

import flyte
import flyte.report

env = flyte.TaskEnvironment(
    name="duckdb-etl",
    image=flyte.Image.from_debian_base().with_pip_packages("duckdb", "pandas", "pyarrow"),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

SAMPLE_CSV = "https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv"

@env.task
async def extract() -> pd.DataFrame:
    """Load raw data from a CSV source."""
    df = duckdb.sql(f"SELECT * FROM read_csv_auto('{SAMPLE_CSV}')").df()
    print(f"Extracted {len(df)} rows")
    return df

@env.task
async def transform(raw: pd.DataFrame) -> pd.DataFrame:
    """Aggregate survival stats by passenger class using DuckDB SQL."""
    summary = duckdb.sql("""
        SELECT
            Pclass AS passenger_class,
            COUNT(*) AS total,
            SUM(Survived) AS survived,
            ROUND(AVG(Survived) * 100, 1) AS survival_rate,
            ROUND(AVG(Fare), 2) AS avg_fare
        FROM raw
        GROUP BY Pclass
        ORDER BY Pclass
    """).df()
    print(summary.to_string(index=False))
    return summary

@env.task(report=True)
async def pipeline() -> pd.DataFrame:
    """Extract → Transform pipeline."""
    raw = await extract()
    summary = await transform(raw)

    await flyte.report.replace.aio(
        f"<h2>DuckDB ETL Results</h2>"
        f"<p>Processed {len(raw)} rows into {len(summary)} groups</p>"
        f"<h3>Survival by Passenger Class</h3>"
        f"{summary.to_html(index=False)}"
    )
    await flyte.report.flush.aio()

    return summary

#flyte run duckdb_etl.py pipeline