Update app.py

app.py

@@ -1,43 +1,63 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline
 from diffusers import DiffusionPipeline
 import torch
+from PIL import Image
+
+device = "cpu" 
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.mps.is_available():
+    device = "mps"
+    
+model_id_image = "CompVis/stable-diffusion-v1-4"
+model_id_image_description = "vikhyatk/moondream2"
+revision = "2024-08-26"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_id = "CompVis/stable-diffusion-v1-4"
 torch_dtype = torch.float32
 
 if torch.cuda.is_available():
     torch_dtype = torch.bfloat16
 
 def generate_description(image):
-    model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-    return model(image)[0]['generated_text']
+    model = AutoModelForCausalLM.from_pretrained(model_id_image_description, trust_remote_code=True, revision=revision)
+    tokenizer = AutoTokenizer.from_pretrained(model_id_image_description, revision=revision)
 
-def generate_image_by_description(description):
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+    image_test = Image.open(image)
+    enc_image = model.encode_image(image_test)
+    res = model.answer_question(enc_image, "Describe this image to create an avatar", tokenizer)
+    return res
+
+def generate_image_by_description(description, avatar_style=None):
+    pipe = StableDiffusionPipeline.from_pretrained(model_id_image, torch_dtype=torch_dtype)
     pipe = pipe.to(device)
     pipe.enable_attention_slicing()
-
+    
     prompt = (
-        f"Generate a image of a pigeon for my profile avatar. "
-        f"The description of the pigeon is: {description}. "
+        f"Create a pigeon profile avatar. "
+        f"Use the following description: {description}. "
         )
+
+    if avatar_style:
+        prompt += f"Use {avatar_style} avatar style."
+
     image = pipe(prompt).images[0]  
     return image
 
 
+def process_and_generate(image, avatar_style):
+    description = generate_description(image)
+    return generate_image_by_description(description, avatar_style)
+
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=2, min_width=300):
             selected_image = gr.Image(type="filepath", label="Upload an Image of the Pigeon",height=300)
+            avatar_style = gr.Radio(
+            ["Realistic", "Pixel Art", "Imaginative", "Cartoon"], label="(optional) Select the avatar style:")
             generate_button = gr.Button("Generate Avatar", variant="primary")
         with gr.Column(scale=2, min_width=300):
             generated_image = gr.Image(type="numpy", label="Generated Avatar", height=300)
-            def process_and_generate(image):
-                description = generate_description(image)
-                return generate_image_by_description(description)
-
-            generate_button.click(process_and_generate, inputs=selected_image, outputs=generated_image)
+            generate_button.click(process_and_generate, inputs=[selected_image, avatar_style ], outputs=generated_image)
 demo.launch()