chore: fix dtype, etc

.gitignore

@@ -0,0 +1,3 @@
+.env
+__pycache__
+.DS_Store

app.py

@@ -1,19 +1,20 @@
+import spaces # import first
+
 import random
 
 import numpy as np
 import torch
 
 from diffusers import StableDiffusionXLPipeline
-import spaces
 
 import gradio as gr
 
-from .tkg import apply_tkg_noise, COLOR_SET_MAP
+from tkg import apply_tkg_noise, ColorSet, COLOR_SET_MAP
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda"
 model_repo_id = "cagliostrolab/animagine-xl-4.0"  # Replace to the model you would like to use
 
 
@@ -44,7 +45,7 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     latents = torch.randn(
         (
@@ -82,7 +83,7 @@ def infer(
 
 def color_name_to_channels(color_name: str) -> list[int]:
     if color_name in COLOR_SET_MAP:
-        return COLOR_SET_MAP[color_name]
+        return COLOR_SET_MAP[color_name].channels
     else:
         raise ValueError(f"Unknown color name: {color_name}")
 
@@ -122,7 +123,15 @@ examples = [
 
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("# TKG Chroma-Key with AnimagineXL 4.0")
+        gr.Markdown(
+            """
+# TKG Chroma-Key with AnimagineXL 4.0
+
+TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
+- arXiv: https://arxiv.org/abs/2411.15580
+- GitHub: https://github.com/ryugo417/TKG-DM
+
+""")
 
         with gr.Row():
             prompt = gr.Text(
@@ -144,7 +153,7 @@ with gr.Blocks() as demo:
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
-                default="lowres, bad anatomy, bad hands, text, error, missing finger, extra digits, fewer digits, cropped, worst quality, low quality, low score, bad score, average score, signature, watermark, username, blurry",
+                value="lowres, bad anatomy, bad hands, text, error, missing finger, extra digits, fewer digits, cropped, worst quality, low quality, low score, bad score, average score, signature, watermark, username, blurry",
             )
 
             seed = gr.Slider(

requirements.txt

@@ -3,4 +3,4 @@ diffusers
 torch
 transformers
 xformers
-spaces
+spaces>=0.40.1