From fe7c3884951e5bff58e22e017bcbfa806684df06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20P=C4=99=C5=9Bko?= <mpesko3031@gmail.com>
Date: Sun, 18 Nov 2018 10:58:21 +0100
Subject: [PATCH] Fix problem with CartoonGAN not change model and add
 timestamp to Comic

---
 api/migrations/0003_add_timestamp.py | 18 ++++++++++++++++++
 api/models.py                        |  3 +--
 api/views.py                         | 16 ++++++++--------
 style_transfer/style_transfer.py     | 12 ++++++++++--
 4 files changed, 37 insertions(+), 12 deletions(-)
 create mode 100644 api/migrations/0003_add_timestamp.py

diff --git a/api/migrations/0003_add_timestamp.py b/api/migrations/0003_add_timestamp.py
new file mode 100644
index 0000000..ebe5ab7
--- /dev/null
+++ b/api/migrations/0003_add_timestamp.py
@@ -0,0 +1,18 @@
+# Generated by Django 2.0.7 on 2018-11-18 10:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('api', '0002_add_additional_info_to_Comic'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='comic',
+            name='timestamp',
+            field=models.DateTimeField(auto_now_add=True),
+        ),
+    ]
diff --git a/api/models.py b/api/models.py
index 2f940a0..bad1e35 100644
--- a/api/models.py
+++ b/api/models.py
@@ -72,13 +72,12 @@ class Comic(models.Model):
     rl_mode = models.PositiveIntegerField(default=0)
     image_assessment_mode = models.PositiveIntegerField(default=0)
     style_transfer_mode = models.PositiveIntegerField(default=0)
+    timestamp = models.DateTimeField(auto_now_add=True)
 
     @classmethod
     @profile
     def create_from_nparray(cls, nparray, video, yt_url, frames_mode,
                             rl_mode, image_assessment_mode, style_transfer_mode):
-        # if nparray.max() <= 1:
-        #     nparray = nparray.astype(int)
         tmp_name = uuid.uuid4().hex + ".png"
         cv2.imwrite(jj(settings.TMP_DIR, tmp_name), nparray)
         with open(jj(settings.TMP_DIR, tmp_name), mode="rb") as tmp_file:
diff --git a/api/views.py b/api/views.py
index 3e5e6e0..bda67c5 100644
--- a/api/views.py
+++ b/api/views.py
@@ -55,18 +55,18 @@ class ComixifyFromYoutube(APIView):
         image_assessment_mode = serializer.validated_data["image_assessment_mode"]
         style_transfer_mode = serializer.validated_data["style_transfer_mode"]
 
-        comix = Comic.objects.filter(yt_url=yt_url,
-                                     frames_mode=frames_mode,
-                                     rl_mode=rl_mode,
-                                     image_assessment_mode=image_assessment_mode,
-                                     style_transfer_mode=style_transfer_mode
-                                     ).first()
-        if comix is not None:
+        try:
+            comix = Comic.objects.filter(yt_url=yt_url,
+                                         frames_mode=frames_mode,
+                                         rl_mode=rl_mode,
+                                         image_assessment_mode=image_assessment_mode,
+                                         style_transfer_mode=style_transfer_mode
+                                         ).latest('timestamp')
             response = {
                 "status_message": "ok",
                 "comic": comix.file.url,
             }
-        else:
+        except Comic.DoesNotExist:
             video = Video()
             _, yt_download_time = video.download_from_youtube(yt_url)
             video.save()
diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py
index f45c896..d116dfe 100644
--- a/style_transfer/style_transfer.py
+++ b/style_transfer/style_transfer.py
@@ -62,8 +62,16 @@ class StyleTransfer():
 
     @classmethod
     def _cartoon_gan_stylize(cls, frames, gpu=True, style='Hayao'):
-        model_cache_key = 'model_cache'
-        model = cache.get(model_cache_key)  # get model from cache
+        if style == 'Hayao':
+            model_cache_key = 'model_cache_hayao'
+            model = cache.get(model_cache_key)  # get model from cache
+
+        elif style == 'Hosoda':
+            model_cache_key = 'model_cache_hosoda'
+            model = cache.get(model_cache_key)  # get model from cache
+
+        else:
+            raise Exception('No such CartoonGAN model!')
 
         if model is None:
             # load pretrained model