Spaces:
No application file
No application file
File size: 3,028 Bytes
bce4c09 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import modal
import io
import os
from services.session_id_generator import SessionIdGenerator
# Define the Modal image with required dependencies
image = modal.Image.debian_slim().pip_install(
"diffusers",
"transformers",
"torch",
"safetensors",
"accelerate",
"Pillow",
"python-dotenv", # Add this for .env support
)
# Load environment variables only when running locally (not in Modal's cloud)
if modal.is_local():
from dotenv import load_dotenv
load_dotenv()
# Set Modal credentials from .env file
modal_token_id = os.environ.get("MODAL_TOKEN_ID")
modal_token_secret = os.environ.get("MODAL_TOKEN_SECRET")
if modal_token_id and modal_token_secret:
os.environ["MODAL_TOKEN_ID"] = modal_token_id
os.environ["MODAL_TOKEN_SECRET"] = modal_token_secret
app = modal.App("comic-image-generator", image=image)
class ComicImageGenerator:
def __init__(self):
from diffusers import AutoPipelineForText2Image
import torch
self.pipe = AutoPipelineForText2Image.from_pretrained(
"stabilityai/sdxl-turbo",
torch_dtype=torch.float16,
variant="fp16",
)
self.pipe.to("cuda")
self.torch = torch
def generate_comic_panel(
self,
prompt: str,
panel_id: int,
session_id: str,
steps: int = 1,
seed: int = 42,
) -> tuple:
import time
generator = self.torch.manual_seed(seed)
start = time.time()
result = self.pipe(
prompt=prompt,
generator=generator,
num_inference_steps=steps,
guidance_scale=0.0,
width=512,
height=512,
output_type="pil",
)
duration = time.time() - start
image = result.images[0]
buf = io.BytesIO()
image.save(buf, format="PNG")
buf.seek(0)
return buf.read(), duration
@app.function(image=image, gpu="A10G")
def generate_comic_panel(
prompt: str,
panel_id: int,
session_id: str,
steps: int = 1,
seed: int = 42,
) -> tuple:
generator = ComicImageGenerator()
return generator.generate_comic_panel(
prompt, panel_id, session_id, steps, seed
)
@app.local_entrypoint()
def main():
prompt = "A K-pop idol walking through a rainy Seoul street, whimsical, soft lighting, watercolor style"
panel_id = 1
session_id = SessionIdGenerator.create_session_id("test")
steps = 1 # SDXL Turbo works best with 1 step
seed = 42
# with app.run():
img_bytes, duration = generate_comic_panel.remote(
prompt, panel_id, session_id, steps, seed
)
# Save to local directory
out_dir = f"storyboard/{session_id}/content"
os.makedirs(out_dir, exist_ok=True)
out_path = f"{out_dir}/panel_{panel_id}.png"
with open(out_path, "wb") as f:
f.write(img_bytes)
print("✅ Image generated at:", out_path)
print("🕒 Time taken:", duration, "seconds")
|