Add app

README.md

@@ -1,13 +1,13 @@
 ---
 title: StableMaterials
-emoji: 🦀
+emoji: 🧱
 colorFrom: blue
 colorTo: blue
 sdk: gradio
-sdk_version: 4.32.2
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 license: openrail
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,57 @@
+import gradio as gr
+import spaces
+import torch
+import diffusers
+
+from generation import generate_material
+
+
+@spaces.GPU
+def generate(prompts, seed, resolution, refinement):
+  image = generate_material(prompts, seed=seed, resolution=int(resolution), refinement=refinement)
+  return image.basecolor, image.normal, image.height, image.metallic, image.roughness
+
+
+def interface_function(prompt_type, text_prompt, image_prompt, seed, resolution, refinement):
+    if prompt_type == "Text":
+        return generate(text_prompt, seed, resolution, refinement)
+    elif prompt_type == "Image":
+        return generate(image_prompt, seed, resolution, refinement)
+
+def update_visibility(prompt_type):
+    if prompt_type == "Text":
+        return gr.update(visible=True), gr.update(visible=False)
+    elif prompt_type == "Image":
+        return gr.update(visible=False), gr.update(visible=True)
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            prompt_type = gr.Radio(choices=["Text", "Image"], label="Prompt Type", value="Text")
+            text_prompt = gr.Textbox(label="Text Prompt", visible=True, lines=3, placeholder="A brick wall")
+            image_prompt = gr.Image(type="pil", label="Image Prompt", visible=False)
+        
+        with gr.Column():
+            seed = gr.Number(value=-1, label="Seed (-1 for random)")
+            resolution = gr.Dropdown(["512", "1024", "2048"], value="512", label="Resolution", interactive=False)
+            refinement = gr.Checkbox(label="Refinement", interactive=False)
+            generate_button = gr.Button("Generate")
+
+    prompt_type.change(fn=update_visibility, inputs=prompt_type, outputs=[text_prompt, image_prompt])
+    
+    with gr.Row():
+        output_basecolor = gr.Image(label="Base Color", format="png", image_mode="RGB")
+        output_normal = gr.Image(label="Normal Map", format="png", image_mode="RGB")
+        output_height = gr.Image(label="Height Map", format="png", image_mode="L")
+        output_metallic = gr.Image(label="Metallic Map", format="png", image_mode="L")
+        output_roughness = gr.Image(label="Roughness Map", format="png", image_mode="L")
+    
+    generate_button.click(
+        fn=interface_function, 
+        inputs=[prompt_type, text_prompt, image_prompt, seed, resolution, refinement], 
+        outputs=[output_basecolor, output_normal, output_height, output_metallic, output_roughness]
+    )
+
+if __name__ == "__main__":
+    demo.launch()
+

generation.py

@@ -0,0 +1,51 @@
+import torch
+from diffusers import DiffusionPipeline, LCMScheduler, UNet2DConditionModel
+
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+unet = UNet2DConditionModel.from_pretrained(
+    "gvecchio/StableMaterials",
+    subfolder="unet_lcm",
+    torch_dtype=torch.float16,
+)
+
+pipe = DiffusionPipeline.from_pretrained(
+    "gvecchio/StableMaterials",
+    trust_remote_code=True,
+    unet=unet,
+    torch_dtype=torch.float16
+).to(device)
+
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+
+
+def generate_material(prompt, seed=-1, resolution=512, refinement=False):
+    try:
+        seed = seed if seed != -1 else torch.randint(0, 10000, (1,)).item()
+        
+        logger.info(f"Generating images for prompt: {prompt} with seed: {seed}")
+        generator = torch.Generator(device=pipe.device).manual_seed(seed)
+
+        image = pipe(
+            prompt=[prompt],
+            tileable=True,
+            num_images_per_prompt=1,
+            num_inference_steps=4,
+            generator=generator,
+        ).images[0]
+
+        image = image.resize((resolution, resolution))
+
+        if refinement:
+            pass
+
+        return image
+    except Exception as e:
+        logger.error(f"Exception occurred while generating images: {e}")
+        raise

requirements.txt

@@ -0,0 +1,6 @@
+accelerate
+diffusers
+invisible_watermark
+torch
+torchvision
+transformers