Update app.py

app.py

@@ -18,13 +18,13 @@ def generate(input_ids,
              cfg_weight: float = 5,
              image_token_num_per_image: int = 576,
              patch_size: int = 16):
-    tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int)  #.cuda()
+    tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).to(cuda_device)
     for i in range(parallel_size * 2):
         tokens[i, :] = input_ids
         if i % 2 != 0:
             tokens[i, 1:-1] = processor.pad_id
     inputs_embeds = model.language_model.get_input_embeddings()(tokens)
-    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int)  #.cuda()
+    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).to(cuda_device)
 
     pkv = None
     for i in range(image_token_num_per_image):
@@ -89,8 +89,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label='Prompt', value='portrait, color, cinematic')
-            width = gr.Slider(256, 1536, 256, step=16, label='Width')
-            height = gr.Slider(256, 1536, 256, step=16, label='Height')
+            width = gr.Slider(128, 1536, 128, step=16, label='Width')
+            height = gr.Slider(128, 1536, 128, step=16, label='Height')
             guidance = gr.Slider(1.0, 10.0, 5, step=0.1, label='Guidance')
             seed = gr.Number(-1, precision=0, label='Seed (-1 for random)')
 
@@ -113,6 +113,7 @@ with gr.Blocks() as demo:
         )
 
 if __name__ == '__main__':
+    cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
     model_path = 'deepseek-ai/Janus-1.3B'
     processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
     tokenizer = processor.tokenizer
@@ -120,6 +121,8 @@ if __name__ == '__main__':
     config = AutoConfig.from_pretrained(model_path)
     config.language_config._attn_implementation = 'eager'
     model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
-    # model = model.to(torch.bfloat16).cuda()
-    model = model.to(torch.float16)
+    if torch.cuda.is_available():
+        model = model.to(torch.bfloat16).cuda()
+    else:
+        model = model.to(torch.float16)
     demo.launch()