text_to_texture / app.py
takahirox's picture
WIP commit
cdd743a
raw
history blame contribute delete
No virus
1.39 kB
from diffusers import StableDiffusionPipeline
from compel import Compel
import gradio
import torch
model_id = "dream-textures/texture-diffusion"
device = "cuda"
dtype = torch.float16
if torch.cuda.is_available():
pipe = StableDiffusionPipeline.from_pretrained(
model_id, torch_dtype=dtype
).to(device)
else:
pipe = StableDiffusionPipeline.from_pretrained(model_id)
compel_proc = Compel(
tokenizer=pipe.tokenizer,
text_encoder=pipe.text_encoder,
truncate_long_prompts=False,
)
def predict(
prompt: str,
generator: int,
num_inference_steps: int,
strength: float,
guidance_scale: float,
):
generator = torch.manual_seed(generator)
prompt_embeds = compel_proc(prompt)
results = pipe(
prompt_embeds=prompt_embeds,
generator=generator,
guidance_scale=float(guidance_scale),
num_inference_steps=num_inference_steps,
output_type="pil",
strength=float(strength),
)
if len(results.images) > 0:
return results.images[0]
return None
app = gradio.Interface(
fn=predict,
inputs=[
gradio.Textbox("pbr brick wall"), # prompt
gradio.Slider(0, 2147483647, 2159232, step=1), # generator
gradio.Slider(2, 15, 4, step=1), # num_inference_steps
gradio.Slider(0.0, 1.0, 0.5, step=0.01), # strength
gradio.Slider(0.0, 5.0, 0.2, step=0.01), # guidance_scale
],
outputs=gradio.Image(type="pil")
)
app.launch()