Spaces:
Sleeping
Sleeping
File size: 1,970 Bytes
4f48282 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
#img_gen_modal.py
import modal
import random
from datetime import datetime
import random
import io
from config.config import prompts, models # Indirect import
import os
import torch
from huggingface_hub import login
from transformers import AutoTokenizer
# Define the Modal image
image = (
#modal.Image.from_registry("nvidia/cuda:12.2.0-devel-ubuntu22.04", add_python="3.9")
modal.Image.debian_slim(python_version="3.9") # Base image
.apt_install(
"git",
)
.pip_install(
"diffusers",
"transformers",
"torch",
"accelerate",
"gradio>=4.44.1",
"safetensors",
"pillow",
"sentencepiece",
"hf_transfer",
"huggingface_hub[hf_transfer]",
"aria2", # aria2 for ultra-fast parallel downloads
f"git+https://github.com/huggingface/transformers.git"
)
.env(
{
"HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_HOME": "HF_HOME"
}
)
)
# Create a Modal app
app = modal.App("img-gen-modal", image=image)
with image.imports():
import diffusers
import os
import gradio
import torch
import sentencepiece
#flux_model_vol = modal.Volume.from_name("flux-model-vol", create_if_missing=True) # Reference your volume
@app.function(
secrets=[modal.Secret.from_name("huggingface-token")],
#volumes={"/data": flux_model_vol},
gpu="t4",
timeout=600
)
def generate_image():
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
prompt = "A cat holding a sign that says hello world"
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")
generate_image()
|