Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,423 Bytes
7aec2a5 ebcdc6c f82344a ebcdc6c 7aec2a5 a49ecc2 7aec2a5 a49ecc2 7aec2a5 66f492a a49ecc2 66f492a 7aec2a5 |
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 |
import gradio as gr
from diffusers import DiffusionPipeline
import os
import torch
import shutil
import spaces
def find_cuda():
# Check if CUDA_HOME or CUDA_PATH environment variables are set
cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
if cuda_home and os.path.exists(cuda_home):
return cuda_home
# Search for the nvcc executable in the system's PATH
nvcc_path = shutil.which('nvcc')
if nvcc_path:
# Remove the 'bin/nvcc' part to get the CUDA installation path
cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
return cuda_path
return None
cuda_path = find_cuda()
if cuda_path:
print(f"CUDA installation found at: {cuda_path}")
else:
print("CUDA installation not found")
# check if cuda is available
device = "cuda" if torch.cuda.is_available() else "cpu"
# load the pipeline/model
pipeline = DiffusionPipeline.from_pretrained("jadechoghari/mar", trust_remote_code=True, custom_pipeline="jadechoghari/mar")
# function that generates images
@spaces.GPU
def generate_image(seed, num_ar_steps, class_labels, cfg_scale, cfg_schedule):
generated_image = pipeline(
model_type="mar_huge", # using mar_huge
seed=seed,
num_ar_steps=num_ar_steps,
class_labels=[int(label.strip()) for label in class_labels.split(',')],
cfg_scale=cfg_scale,
cfg_schedule=cfg_schedule,
output_dir="./images"
)
return generated_image
with gr.Blocks() as demo:
gr.Markdown("""
# MAR Image Generation Demo 🚀
Welcome to the demo for **MAR** (Masked Autoregressive Model), a novel approach to image generation that eliminates the need for vector quantization. MAR uses a diffusion process to generate images in a continuous-valued space, resulting in faster, more efficient, and higher-quality outputs.
Simply adjust the parameters below to create your custom images in real-time.
Make sure to provide valid **ImageNet class labels** to see the translation of text to image. For a complete list of ImageNet classes, check out [this reference](https://deeplearning.cms.waikato.ac.nz/user-guide/class-maps/IMAGENET/).
For more details, visit the [GitHub repository](https://github.com/LTH14/mar).
""")
seed = gr.Number(value=0, label="Seed")
num_ar_steps = gr.Slider(minimum=1, maximum=256, value=64, label="Number of AR Steps")
class_labels = gr.Textbox(value="207, 360, 388, 113, 355, 980, 323, 979", label="Class Labels (comma-separated ImageNet labels)")
cfg_scale = gr.Slider(minimum=1, maximum=10, value=4, label="CFG Scale")
cfg_schedule = gr.Dropdown(choices=["constant", "linear"], label="CFG Schedule", value="constant")
image_output = gr.Image(label="Generated Image")
generate_button = gr.Button("Generate Image")
# we link the button to the function and display the output
generate_button.click(generate_image, inputs=[seed, num_ar_steps, class_labels, cfg_scale, cfg_schedule], outputs=image_output)
# Add example inputs and cache them for quicker demos
gr.Interface(
generate_image,
inputs=[seed, num_ar_steps, class_labels, cfg_scale, cfg_schedule],
outputs=image_output,
# examples=[
# [0, 64, "207, 360, 388, 113, 355, 980, 323, 979", 4, "constant"],
# ],
# cache_examples=True
)
demo.launch()
|