Update app.py

app.py

@@ -1,51 +1,54 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, AnimateDiffPipeline
-from PIL import Image
+from diffusers import StableDiffusionPipeline
+from transformers import AutoProcessor, AutoModel
+import os
 
-# Load the 3D Cartoon Image Generator (DreamShaper)
+# Load Text-to-Image model (Redshift Diffusion)
 image_pipe = StableDiffusionPipeline.from_pretrained(
-    "Lykon/DreamShaper", torch_dtype=torch.float16
-).to("cuda")
+    "nitrosocke/redshift-diffusion", torch_dtype=torch.float16
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load the Animation Model (AnimateDiff)
-video_pipe = AnimateDiffPipeline.from_pretrained(
-    "guoyww/animatediff", torch_dtype=torch.float16
-).to("cuda")
+# Load Image-to-Video model (Zeroscope v2 XL)
+video_model_id = "cerspense/zeroscope_v2_XL"
+processor = AutoProcessor.from_pretrained(video_model_id)
+video_model = AutoModel.from_pretrained(video_model_id).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Function to generate a 3D cartoon image
-def generate_cartoon_image(prompt):
+# Function to generate image from text
+def generate_image(prompt):
     image = image_pipe(prompt).images[0]
-    image_path = "generated_cartoon.png"
+    image_path = "generated_image.png"
     image.save(image_path)
     return image_path
 
-# Function to generate an animated video from an uploaded image
-def generate_cartoon_video(image):
-    image = Image.open(image).convert("RGB")
-    video_output = video_pipe(image)["video"]
-    video_path = "generated_cartoon_animation.mp4"
+# Function to convert image to video
+def generate_video(image_path):
+    with torch.no_grad():
+        inputs = processor(images=image_path, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+        video_output = video_model(**inputs)
+    
+    video_path = "generated_video.mp4"
     video_output.save(video_path)
     return video_path
 
-# Gradio UI
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 AI 3D Cartoon Image & Animation Generator")
+    gr.Markdown("## 🎨 AI Cartoon Image & Video Generator")
     
-    # Text-to-Image Section
-    with gr.Tab("Generate Cartoon Image"):
-        gr.Markdown("Enter a text prompt to generate a 3D cartoon-style image")
-        text_input = gr.Textbox(label="Enter your prompt")
-        img_output = gr.Image(label="Generated Cartoon Image")
-        img_button = gr.Button("Generate Image")
-        img_button.click(generate_cartoon_image, inputs=text_input, outputs=img_output)
-
-    # Image-to-Video Section
-    with gr.Tab("Generate Cartoon Animation"):
-        gr.Markdown("Upload a 3D cartoon-style image to create an animation")
-        img_input = gr.Image(type="filepath", label="Upload Cartoon Image")
-        vid_output = gr.Video(label="Generated Animation")
-        vid_button = gr.Button("Generate Animation")
-        vid_button.click(generate_cartoon_video, inputs=img_input, outputs=vid_output)
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter Text Prompt", placeholder="A 3D cartoon cat playing in a park")
+        generate_image_btn = gr.Button("Generate Image")
+    
+    image_output = gr.Image(label="Generated Image")
+    
+    with gr.Row():
+        generate_video_btn = gr.Button("Convert to Video")
+        video_output = gr.Video(label="Generated Video")
+    
+    download_image = gr.File(label="Download Image")
+    download_video = gr.File(label="Download Video")
+    
+    generate_image_btn.click(generate_image, inputs=[prompt_input], outputs=[image_output, download_image])
+    generate_video_btn.click(generate_video, inputs=[image_output], outputs=[video_output, download_video])
 
 demo.launch()