acek / app.py
acek99's picture
Upload 2 files
bf4de7d verified
Raw History Blame Contribute Delete
3.35 kB
"""Bang Story image Space: draws clip pictures on Hugging Face ZeroGPU.
Upload this file together with requirements.txt and README.md into a Gradio Space, set the Space's
hardware to ZeroGPU, and Bang Story can use it as its image service — free, with no credits and no
API keys. The app calls the `generate` function through the Space's Gradio API.
"""
import os
import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
# Tried in order, first one that loads wins. The official repository gives the best quality but is gated:
# without the licence accepted and a read token in the HF_TOKEN secret it answers 401. The open mirrors
# below are the same model, so the Space still comes up instead of dying on a missing token.
MODEL_CHAIN = [
os.environ.get("MODEL_ID") or "black-forest-labs/FLUX.1-schnell",
"Niansuh/FLUX.1-schnell",
"playgroundai/playground-v2.5-1024px-aesthetic",
]
# A distilled model like schnell needs very few steps; a full pipeline needs many more.
MAX_STEPS = 8
device = "cuda" if torch.cuda.is_available() else "cpu"
def load():
"""Loads the first model of the chain that works, so one blocked repository cannot kill the Space."""
last = None
for model in MODEL_CHAIN:
if not model:
continue
try:
print(f"[bang-story] loading {model} ...", flush=True)
pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.bfloat16).to(device)
print(f"[bang-story] ready: {model} on {device}", flush=True)
return pipe, model
except Exception as e:
last = e
print(f"[bang-story] {model} failed: {type(e).__name__}: {e}", flush=True)
raise RuntimeError(f"Tidak ada model yang bisa dimuat. Error terakhir: {last}")
pipe, MODEL = load()
IS_FLUX = "flux" in MODEL.lower()
@spaces.GPU()
def generate(prompt: str, width: int = 1024, height: int = 576, steps: int = 4):
"""Draws one picture. Bang Story calls this through the Gradio API."""
kwargs = {"prompt": prompt, "width": int(width), "height": int(height)}
if IS_FLUX:
kwargs["num_inference_steps"] = max(1, min(MAX_STEPS, int(steps)))
# schnell is distilled: it runs without classifier-free guidance.
kwargs["guidance_scale"] = 0.0
else:
kwargs["num_inference_steps"] = max(20, min(35, int(steps) * 5))
kwargs["guidance_scale"] = 7.0
return pipe(**kwargs).images[0]
with gr.Blocks(title="Bang Story — Image") as demo:
gr.Markdown(
f"Draws the clip pictures for **Bang Story** with `{MODEL}`.\n\n"
"Bang Story calls this Space from its Settings, so you normally never touch this page. "
"It is left open so you can try a prompt yourself."
)
prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Describe the picture...")
with gr.Row():
width = gr.Slider(256, 1536, value=1024, step=16, label="Width")
height = gr.Slider(256, 1536, value=576, step=16, label="Height")
steps = gr.Slider(1, MAX_STEPS, value=4, step=1, label="Steps")
run = gr.Button("Generate", variant="primary")
out = gr.Image(label="Result", type="pil")
run.click(generate, [prompt, width, height, steps], out)
prompt.submit(generate, [prompt, width, height, steps], out)
demo.launch()