Update app.py

app.py

@@ -2,9 +2,18 @@ import torch
 import gradio as gr
 from diffusers import AnimateDiffPipeline, MotionAdapter, DPMSolverMultistepScheduler, AutoencoderKL, SparseControlNetModel, EulerAncestralDiscreteScheduler
 from diffusers.utils import export_to_gif, load_image
+from realesrgan import RealESRGAN
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+def enhance_quality(image_path):
+    model = RealESRGAN(device, scale=4)
+    model.load_weights('RealESRGAN_x4.pth', download=True)
+    img = Image.open(image_path)
+    sr_image = model.predict(img)
+    sr_image.save('enhanced_' + image_path)
+    return 'enhanced_' + image_path
+
 def generate_video(prompt, negative_prompt, num_inference_steps, conditioning_frame_indices, controlnet_conditioning_scale):
     motion_adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-3", torch_dtype=torch.float16).to(device)
     controlnet = SparseControlNetModel.from_pretrained("guoyww/animatediff-sparsectrl-scribble", torch_dtype=torch.float16).to(device)
@@ -41,7 +50,8 @@ def generate_video(prompt, negative_prompt, num_inference_steps, conditioning_fr
     ).frames[0]
     
     export_to_gif(video, "output.gif")
-    return "output.gif"
+    enhanced_gif = enhance_quality("output.gif")
+    return enhanced_gif
 
 def generate_simple_video(prompt):
     adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16).to(device)
@@ -65,7 +75,8 @@ def generate_simple_video(prompt):
     ).frames[0]
     
     export_to_gif(frames, "simple_output.gif")
-    return "simple_output.gif"
+    enhanced_gif = enhance_quality("simple_output.gif")
+    return enhanced_gif
 
 demo1 = gr.Interface(
     fn=generate_video,
@@ -91,5 +102,6 @@ demo2 = gr.Interface(
 
 demo = gr.TabbedInterface([demo1, demo2], ["Advanced Video Generation", "Simple Video Generation"])
 
+
 demo.launch()
 #demo.launch(server_name="0.0.0.0", server_port=7910)