update app.py

app.py

@@ -1,17 +1,10 @@
 import gradio as gr
-import cv2
 from diffusers import StableDiffusionPipeline, UNet2DConditionModel
 import torch
 import random
 import numpy as np
 
-
-MODEL="/mnt/bn/ailab-yuningshen-psg/mlx/users/quanquan.gu/playground/trl/iter0_0_2.0e-5_linear200_new/checkpoints/checkpoint_0"
-MODEL="/mnt/bn/ailab-yuningshen-psg/mlx/users/quanquan.gu/playground/trl/iter2_12h_5.0e-8_beta5_rep_wtie/checkpoints/checkpoint_10"
-PROMPTS="/mnt/bn/ailab-yuningshen-psg/mlx/users/quanquan.gu/playground/trl/dataset/pickapic/pickapic_v2/validation_unique-00007_filtered.parquet"
-NUM=1
-
-
+MODEL="UCLA-AGI/SPIN-Diffusion-iter3"
 
 def set_seed(seed=5775709):
     random.seed(seed)
@@ -19,6 +12,7 @@ def set_seed(seed=5775709):
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
 
+
 set_seed()
 
 def get_pipeline(device='cuda'):
@@ -43,11 +37,6 @@ def generate(prompt: str, num_images: int=5, guidance_scale=7.5):
     images = [x.resize((512, 512)) for x in images]
     return images
 
-def gen_image(args, image):
-    output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    return output
-
-
 
 
 with gr.Blocks() as demo:
@@ -76,4 +65,4 @@ with gr.Blocks() as demo:
     generate_btn.click(fn=generate, inputs=[prompt_input, num_images_input, guidance_scale], outputs=gallery)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)